""" Data loading and preprocessing module for chromatographic RT prediction. Handles SMILES string to molecular graph conversion and feature extraction. """ import pandas as pd import numpy as np from rdkit import Chem, DataStructs from rdkit.Chem import Descriptors, Crippen, Lipinski, rdMolDescriptors, rdFingerprintGenerator from sklearn.preprocessing import LabelEncoder from sklearn.model_selection import KFold, StratifiedKFold import torch from torch_geometric.data import Data, DataLoader from typing import List, Tuple, Dict, Optional, Any import warnings warnings.filterwarnings('ignore') class MolecularFeatureExtractor: """Extract molecular features from SMILES strings.""" def __init__(self): self.atom_features = [ 'atomic_num', 'degree', 'formal_charge', 'hybridization', 'is_aromatic', 'total_numHs', 'mass', 'chiral_tag' ] self.bond_features = ['bond_type', 'conjugated', 'is_ring', 'stereo'] def get_atom_features(self, atom) -> List[float]: """Extract atom-level features including chirality.""" chiral_tag = atom.GetChiralTag() # Convert chiral tag to numerical value chiral_value = { Chem.ChiralType.CHI_UNSPECIFIED: 0.0, Chem.ChiralType.CHI_TETRAHEDRAL_CW: 1.0, # R Chem.ChiralType.CHI_TETRAHEDRAL_CCW: 2.0, # S Chem.ChiralType.CHI_OTHER: 3.0, Chem.ChiralType.CHI_TETRAHEDRAL: 4.0, Chem.ChiralType.CHI_ALLENE: 5.0, Chem.ChiralType.CHI_SQUAREPLANAR: 6.0, Chem.ChiralType.CHI_TRIGONALBIPYRAMIDAL: 7.0, Chem.ChiralType.CHI_OCTAHEDRAL: 8.0 }.get(chiral_tag, 0.0) return [ atom.GetAtomicNum(), atom.GetTotalDegree(), atom.GetFormalCharge(), int(atom.GetHybridization()), int(atom.GetIsAromatic()), atom.GetTotalNumHs(), atom.GetMass() / 100.0, # Normalized mass chiral_value / 8.0 # Normalized chirality ] def get_bond_features(self, bond) -> List[float]: """Extract bond-level features including stereochemistry.""" stereo = bond.GetStereo() # Convert stereo to numerical value stereo_value = { Chem.BondStereo.STEREONONE: 0.0, Chem.BondStereo.STEREOANY: 1.0, Chem.BondStereo.STEREOZ: 2.0, # Z (cis) Chem.BondStereo.STEREOE: 3.0, # E (trans) Chem.BondStereo.STEREOCIS: 4.0, Chem.BondStereo.STEREOTRANS: 5.0 }.get(stereo, 0.0) return [ int(bond.GetBondType()), int(bond.GetIsConjugated()), int(bond.IsInRing()), stereo_value / 5.0 # Normalized stereochemistry ] def smiles_to_graph(self, smiles: str) -> Optional[Data]: """Convert SMILES string to PyTorch Geometric graph with stereochemistry.""" try: mol = Chem.MolFromSmiles(smiles) if mol is None: return None # Try to assign stereochemistry from SMILES try: Chem.AssignStereochemistry(mol, cleanIt=True, force=True, flagPossibleStereoCenters=True) except: pass # Continue even if stereochemistry assignment fails # Atom features atom_features = [] for atom in mol.GetAtoms(): atom_features.append(self.get_atom_features(atom)) # Bond features and edge indices edge_indices = [] edge_attrs = [] for bond in mol.GetBonds(): i = bond.GetBeginAtomIdx() j = bond.GetEndAtomIdx() edge_indices.extend([[i, j], [j, i]]) # Undirected graph bond_feat = self.get_bond_features(bond) edge_attrs.extend([bond_feat, bond_feat]) # Convert to tensors x = torch.tensor(atom_features, dtype=torch.float) edge_index = torch.tensor(edge_indices, dtype=torch.long).t().contiguous() if edge_indices else torch.empty((2, 0), dtype=torch.long) edge_attr = torch.tensor(edge_attrs, dtype=torch.float) if edge_attrs else None return Data(x=x, edge_index=edge_index, edge_attr=edge_attr) except Exception as e: print(f"Error processing SMILES {smiles}: {e}") return None def get_molecular_descriptors(self, smiles: str) -> Dict[str, float]: """Extract essential molecular descriptors (simplified for speed).""" try: mol = Chem.MolFromSmiles(smiles) if mol is None: return {} descriptors = { # Basic descriptors - fast to compute 'MW': Descriptors.MolWt(mol), # type: ignore[attr-defined] 'LogP': Crippen.MolLogP(mol), # type: ignore[attr-defined] 'NumHDonors': Lipinski.NumHDonors(mol), # type: ignore[attr-defined] 'NumHAcceptors': Lipinski.NumHAcceptors(mol), # type: ignore[attr-defined] 'NumRotatableBonds': Descriptors.NumRotatableBonds(mol), # type: ignore[attr-defined] 'TPSA': Descriptors.TPSA(mol), # type: ignore[attr-defined] 'NumAromaticRings': Descriptors.NumAromaticRings(mol), # type: ignore[attr-defined] 'NumAliphaticRings': Descriptors.NumAliphaticRings(mol), # type: ignore[attr-defined] 'HeavyAtomCount': mol.GetNumHeavyAtoms(), 'NumHeteroatoms': Descriptors.NumHeteroatoms(mol), # type: ignore[attr-defined] 'RingCount': Descriptors.RingCount(mol), # type: ignore[attr-defined] 'BertzCT': Descriptors.BertzCT(mol), # type: ignore[attr-defined] 'MolMR': Crippen.MolMR(mol), # type: ignore[attr-defined] 'LabuteASA': Descriptors.LabuteASA(mol), # type: ignore[attr-defined] 'HallKierAlpha': Descriptors.HallKierAlpha(mol), # type: ignore[attr-defined] 'Kappa1': Descriptors.Kappa1(mol), # type: ignore[attr-defined] 'Kappa2': Descriptors.Kappa2(mol), # type: ignore[attr-defined] 'Kappa3': Descriptors.Kappa3(mol), # type: ignore[attr-defined] } # Add optional descriptors with error handling try: descriptors['FractionCsp3'] = Descriptors.FractionCsp3(mol) # type: ignore[attr-defined] except: descriptors['FractionCsp3'] = 0.0 # Add simple calculated features descriptors['HeavyAtomRatio'] = descriptors['HeavyAtomCount'] / max(mol.GetNumAtoms(), 1) descriptors['AromaticRatio'] = descriptors['NumAromaticRings'] / max(descriptors['RingCount'], 1) descriptors['RotatableBondRatio'] = descriptors['NumRotatableBonds'] / max(descriptors['HeavyAtomCount'], 1) return descriptors except Exception as e: print(f"Error calculating descriptors for {smiles}: {e}") # Return default values if calculation fails return { 'MW': 200.0, 'LogP': 2.0, 'NumHDonors': 1.0, 'NumHAcceptors': 2.0, 'NumRotatableBonds': 3.0, 'TPSA': 50.0, 'NumAromaticRings': 1.0, 'NumAliphaticRings': 0.0, 'HeavyAtomCount': 15.0, 'NumHeteroatoms': 2.0, 'RingCount': 1.0, 'BertzCT': 100.0, 'MolMR': 60.0, 'LabuteASA': 100.0, 'HallKierAlpha': 5.0, 'Kappa1': 3.0, 'Kappa2': 2.0, 'Kappa3': 1.0, 'FractionCsp3': 0.5, 'HeavyAtomRatio': 0.75, 'AromaticRatio': 1.0, 'RotatableBondRatio': 0.2 } def get_morgan_fingerprint( self, smiles: str, *, n_bits: int = 2048, radius: int = 2, use_chirality: bool = True, use_features: bool = False, ) -> np.ndarray: """Generate a Morgan fingerprint as a dense float32 numpy array.""" try: mol = Chem.MolFromSmiles(smiles) if mol is None: return np.zeros((n_bits,), dtype=np.float32) if use_features: atom_inv_gen = rdFingerprintGenerator.GetMorganFeatureAtomInvGen() else: # Include ring membership information to stay close to RDKit defaults atom_inv_gen = rdFingerprintGenerator.GetMorganAtomInvGen(includeRingMembership=True) generator = rdFingerprintGenerator.GetMorganGenerator( radius=radius, fpSize=n_bits, includeChirality=use_chirality, atomInvariantsGenerator=atom_inv_gen, ) fingerprint = generator.GetFingerprint(mol) array = np.zeros((n_bits,), dtype=np.float32) DataStructs.ConvertToNumpyArray(fingerprint, array) return array except Exception as e: print(f"Error generating fingerprint for {smiles}: {e}") return np.zeros((n_bits,), dtype=np.float32) class RTDataset: """Dataset class for retention time prediction.""" def __init__(self, data_path: str, is_test: bool = False): self.data_path = data_path self.is_test = is_test self.feature_extractor = MolecularFeatureExtractor() self.lab_encoder = None self.valid_indices: Optional[np.ndarray] = None # Load data self.df = pd.read_csv(data_path) print(f"Loaded {'test' if is_test else 'train'} data: {self.df.shape}") def preprocess_data(self, lab_encoder=None) -> Tuple[List[Data], np.ndarray, Optional[np.ndarray]]: """ Preprocess the dataset. Returns: graphs: List of PyTorch Geometric Data objects lab_features: Encoded lab features targets: Target RT values (None for test data) """ # Encode lab features if lab_encoder is None: self.lab_encoder = LabelEncoder() lab_features = self.lab_encoder.fit_transform(self.df['Lab']) else: self.lab_encoder = lab_encoder lab_features = self.lab_encoder.transform(self.df['Lab']) # Convert SMILES to graphs graphs = [] valid_indices = [] print("Converting SMILES to molecular graphs...") for idx, smiles in enumerate(self.df['SMILES']): graph = self.feature_extractor.smiles_to_graph(smiles) if graph is not None: graphs.append(graph) valid_indices.append(idx) else: print(f"Failed to process SMILES at index {idx}: {smiles}") lab_features = np.asarray(lab_features) lab_features = lab_features[valid_indices] targets_array = None if self.is_test else np.asarray(self.df['RT'].values) targets = None if targets_array is None else targets_array[valid_indices] print(f"Successfully processed {len(graphs)} molecules out of {len(self.df)}") self.valid_indices = np.array(valid_indices, dtype=int) return graphs, lab_features, targets def get_molecular_descriptors_df(self) -> pd.DataFrame: """Get traditional molecular descriptors as DataFrame.""" descriptors_list: List[Dict[str, Any]] = [] print("Extracting molecular descriptors...") for idx, smiles in enumerate(self.df['SMILES']): desc = self.feature_extractor.get_molecular_descriptors(smiles) desc_entry: Dict[str, Any] = dict(desc) desc_entry['index'] = idx desc_entry['Lab'] = self.df.loc[idx, 'Lab'] if not self.is_test: desc_entry['RT'] = self.df.loc[idx, 'RT'] descriptors_list.append(desc_entry) desc_df = pd.DataFrame(descriptors_list) # Handle missing values numeric_cols = desc_df.select_dtypes(include=[np.number]).columns desc_df[numeric_cols] = desc_df[numeric_cols].fillna(desc_df[numeric_cols].median()) return desc_df def get_fingerprint_matrix( self, *, n_bits: int = 2048, radius: int = 2, use_chirality: bool = True, use_features: bool = False, ) -> np.ndarray: """Return Morgan fingerprints aligned with the valid graph indices.""" if self.valid_indices is None: raise RuntimeError("Call preprocess_data before requesting fingerprints.") fingerprints: List[np.ndarray] = [] for idx in self.valid_indices: smiles = str(self.df.loc[idx, 'SMILES']) fp = self.feature_extractor.get_morgan_fingerprint( smiles, n_bits=n_bits, radius=radius, use_chirality=use_chirality, use_features=use_features, ) fingerprints.append(fp) if not fingerprints: return np.zeros((0, n_bits), dtype=np.float32) return np.stack(fingerprints, axis=0) def create_cross_validation_splits(n_samples: int, n_splits: int = 5, stratify_col: Optional[np.ndarray] = None, random_state: int = 42) -> List[Tuple[np.ndarray, np.ndarray]]: """Create cross-validation splits.""" if stratify_col is not None: # Use stratified k-fold for categorical stratification kf = StratifiedKFold(n_splits=n_splits, shuffle=True, random_state=random_state) splits = list(kf.split(np.arange(n_samples), stratify_col)) else: # Use regular k-fold kf = KFold(n_splits=n_splits, shuffle=True, random_state=random_state) splits = list(kf.split(np.arange(n_samples))) return splits def create_graph_dataloader(graphs: List[Data], lab_features: np.ndarray, targets: Optional[np.ndarray] = None, indices: Optional[np.ndarray] = None, batch_size: int = 32, shuffle: bool = True): """Create PyTorch Geometric DataLoader.""" if indices is not None: selected_graphs = [graphs[i] for i in indices] selected_lab_features = lab_features[indices] selected_targets = targets[indices] if targets is not None else None else: selected_graphs = graphs selected_lab_features = lab_features selected_targets = targets # Add lab features and targets to graph data for i, graph in enumerate(selected_graphs): graph.lab_feature = torch.tensor([selected_lab_features[i]], dtype=torch.long) if selected_targets is not None: graph.y = torch.tensor([selected_targets[i]], dtype=torch.float) return DataLoader(selected_graphs, batch_size=batch_size, shuffle=shuffle) if __name__ == "__main__": # Test the data loading and preprocessing train_dataset = RTDataset("train.csv", is_test=False) graphs, lab_features, targets = train_dataset.preprocess_data() print(f"Number of graphs: {len(graphs)}") print(f"Lab features shape: {lab_features.shape}") if targets is not None: print(f"Targets shape: {targets.shape}") print(f"Example graph: {graphs[0]}") # Test descriptor extraction desc_df = train_dataset.get_molecular_descriptors_df() print(f"Descriptors shape: {desc_df.shape}") print(f"Descriptors columns: {list(desc_df.columns)}")