| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Dict, Tuple, Optional, List |
|
|
| import numpy as np |
| import pandas as pd |
| from rdkit import Chem, DataStructs, RDLogger |
| from rdkit.Chem import AllChem, Descriptors, Lipinski, Crippen |
|
|
|
|
| @dataclass(frozen=True) |
| class FeatConfig: |
| fp_radius: int = 2 |
| fp_nbits: int = 2048 |
|
|
|
|
| def _morgan_fp(mol, radius: int, nbits: int) -> np.ndarray: |
| fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=nbits) |
| arr = np.zeros((nbits,), dtype=np.int8) |
| DataStructs.ConvertToNumpyArray(fp, arr) |
| return arr |
|
|
|
|
| def _physchem_from_mol(mol) -> Dict[str, float]: |
| return { |
| "rd_mw": float(Descriptors.MolWt(mol)), |
| "rd_logp": float(Crippen.MolLogP(mol)), |
| "rd_tpsa": float(Descriptors.TPSA(mol)), |
| "rd_hbd": float(Lipinski.NumHDonors(mol)), |
| "rd_hba": float(Lipinski.NumHAcceptors(mol)), |
| "rd_rotb": float(Lipinski.NumRotatableBonds(mol)), |
| "rd_rings": float(Lipinski.RingCount(mol)), |
| "rd_heavy_atoms": float(Descriptors.HeavyAtomCount(mol)), |
| } |
|
|
|
|
| def _merge_override(row: pd.Series, rd_vals: Dict[str, float], override_physchem: bool) -> Dict[str, float]: |
| if not override_physchem: |
| return rd_vals |
| out = {} |
| for k, v in rd_vals.items(): |
| if k in row.index and pd.notna(row[k]): |
| try: |
| out[k] = float(row[k]) |
| except Exception: |
| out[k] = v |
| else: |
| out[k] = v |
| return out |
|
|
|
|
| def featurize_smiles( |
| df: pd.DataFrame, |
| smiles_col: str = "smiles", |
| config: FeatConfig = FeatConfig(), |
| add_physchem: bool = True, |
| drop_invalid: bool = True, |
| override_physchem: bool = True, |
| suppress_rdkit_warnings: bool = True, |
| ) -> Tuple[pd.DataFrame, pd.Series]: |
| """ |
| Returns: |
| X (features dataframe), valid_mask (bool Series aligned to input df) |
| |
| Key guarantee: |
| - Internal lists keep same length as df (one entry per row), |
| then we filter by valid_mask at the end (prevents shape mismatch). |
| """ |
| if suppress_rdkit_warnings: |
| RDLogger.DisableLog("rdApp.*") |
|
|
| smiles_list = df[smiles_col].astype(str).tolist() |
| n = len(smiles_list) |
|
|
| |
| dummy = _physchem_from_mol(Chem.MolFromSmiles("CC")) |
| phys_keys = list(dummy.keys()) |
|
|
| fps_all: List[Optional[np.ndarray]] = [None] * n |
| phys_all: List[Optional[Dict[str, float]]] = [None] * n |
| valid: List[bool] = [False] * n |
|
|
| for i, s in enumerate(smiles_list): |
| mol = Chem.MolFromSmiles(s) |
| if mol is None: |
| valid[i] = False |
| continue |
|
|
| valid[i] = True |
| fps_all[i] = _morgan_fp(mol, config.fp_radius, config.fp_nbits) |
|
|
| if add_physchem: |
| rd_vals = _physchem_from_mol(mol) |
| phys_all[i] = _merge_override(df.iloc[i], rd_vals, override_physchem) |
|
|
| valid_mask = pd.Series(valid, index=df.index) |
|
|
| if drop_invalid: |
| keep_idx = df.index[valid_mask] |
| fps_kept = [fp for fp, ok in zip(fps_all, valid) if ok] |
| if len(fps_kept) == 0: |
| raise ValueError("No valid SMILES found after filtering.") |
| fp_mat = np.vstack(fps_kept) |
|
|
| X_fp = pd.DataFrame( |
| fp_mat, |
| columns=[f"fp_{i}" for i in range(config.fp_nbits)], |
| index=keep_idx, |
| ) |
|
|
| if add_physchem: |
| phys_kept = [p for p, ok in zip(phys_all, valid) if ok] |
| X_phys = pd.DataFrame(phys_kept, index=keep_idx) |
| |
| for k in phys_keys: |
| if k not in X_phys.columns: |
| X_phys[k] = np.nan |
| X_phys = X_phys[phys_keys] |
| X = pd.concat([X_phys, X_fp], axis=1) |
| else: |
| X = X_fp |
|
|
| return X, valid_mask |
|
|
| |
| fp_mat = np.vstack([fp if fp is not None else np.zeros((config.fp_nbits,), dtype=np.int8) for fp in fps_all]) |
| X_fp = pd.DataFrame(fp_mat, columns=[f"fp_{i}" for i in range(config.fp_nbits)], index=df.index) |
|
|
| if add_physchem: |
| phys_rows = [] |
| for p in phys_all: |
| if p is None: |
| phys_rows.append({k: np.nan for k in phys_keys}) |
| else: |
| phys_rows.append(p) |
| X_phys = pd.DataFrame(phys_rows, index=df.index) |
| X_phys = X_phys[phys_keys] |
| X = pd.concat([X_phys, X_fp], axis=1) |
| else: |
| X = X_fp |
|
|
| return X, valid_mask |
|
|
