Spaces:
Running
Running
| import argparse | |
| import os | |
| import pandas as pd | |
| import numpy as np | |
| import re | |
| import ast | |
| from rdkit import Chem | |
| from rdkit.Chem import Descriptors | |
| def compute_molecule_stats_numeric(smiles_list, descriptor_funcs, dataset_label, avg_components=np.nan, std_components=np.nan): | |
| """ | |
| Compute statistics (mean, std, min, max) for a given list of SMILES. | |
| Includes average number of components in mixtures (mean and std). | |
| Parameters | |
| ---------- | |
| smiles_list : list | |
| List of SMILES strings. | |
| descriptor_funcs : dict | |
| Dictionary of RDKit descriptor functions. | |
| dataset_label : str | |
| Name of the dataset. | |
| avg_components : float, optional | |
| Average number of components in the mixture. | |
| std_components : float, optional | |
| Standard deviation of number of components in the mixture. | |
| Returns | |
| ------- | |
| list of dict | |
| A single-element list containing the computed statistics. | |
| """ | |
| atom_counts, frag_counts, charges = [], [], [] | |
| descriptor_records = [] | |
| for smi in smiles_list: | |
| mol = Chem.MolFromSmiles(smi) | |
| if mol is not None: | |
| atom_counts.append(mol.GetNumAtoms()) | |
| frag_counts.append(len(Chem.GetMolFrags(mol))) | |
| charges.append(Chem.GetFormalCharge(mol)) | |
| descriptors = {name: func(mol) for name, func in descriptor_funcs.items()} | |
| descriptor_records.append(descriptors) | |
| if atom_counts: | |
| descriptor_df = pd.DataFrame(descriptor_records) | |
| return [{ | |
| 'Dataset Name': dataset_label, | |
| 'Unique molecules': len(atom_counts), | |
| # Atoms | |
| 'Avg atoms/mol': np.mean(atom_counts), | |
| 'Std atoms/mol': np.std(atom_counts), | |
| 'Max atoms/mol': np.max(atom_counts), | |
| 'Min atoms/mol': np.min(atom_counts), | |
| # Fragments | |
| 'Avg fragments': np.mean(frag_counts), | |
| 'Std fragments': np.std(frag_counts), | |
| 'Max fragments': np.max(frag_counts), | |
| # Molecular weight | |
| 'Avg MolWt': descriptor_df['MolWt'].mean(), | |
| 'Std MolWt': descriptor_df['MolWt'].std(), | |
| # Rotatable bonds | |
| 'Avg Rotatable Bonds': descriptor_df['NumRotatableBonds'].mean(), | |
| 'Std Rotatable Bonds': descriptor_df['NumRotatableBonds'].std(), | |
| # Formal charge | |
| 'Avg Formal Charge': np.mean(charges), | |
| 'Std Formal Charge': np.std(charges), | |
| # Components mixture | |
| 'Avg components mixture': avg_components, | |
| 'Std components mixture': std_components | |
| }] | |
| else: | |
| return [{ | |
| 'Dataset Name': dataset_label, | |
| 'Unique molecules': 0, | |
| 'Avg atoms/mol': np.nan, | |
| 'Std atoms/mol': np.nan, | |
| 'Max atoms/mol': np.nan, | |
| 'Min atoms/mol': np.nan, | |
| 'Avg fragments': np.nan, | |
| 'Std fragments': np.nan, | |
| 'Max fragments': np.nan, | |
| 'Avg MolWt': np.nan, | |
| 'Std MolWt': np.nan, | |
| 'Avg Rotatable Bonds': np.nan, | |
| 'Std Rotatable Bonds': np.nan, | |
| 'Avg Formal Charge': np.nan, | |
| 'Std Formal Charge': np.nan, | |
| 'Avg components mixture': avg_components, | |
| 'Std components mixture': std_components | |
| }] | |
| def get_molecule_statistics_with_components(root_path, file_extensions=('.csv',)): | |
| """ | |
| Compute molecular statistics (mean, std, min, max) for each dataset, | |
| including average number of components in mixtures. | |
| Gives preference to processed*.csv files for component counts. | |
| Parameters | |
| ---------- | |
| root_path : str | |
| Path to the root directory containing dataset folders. | |
| file_extensions : tuple of str, optional | |
| File extensions to include in the search (default: '.csv'). | |
| Returns | |
| ------- | |
| pd.DataFrame | |
| Table containing computed molecular statistics and mixture component | |
| information for each dataset or dataset subset (for IlThermo). | |
| """ | |
| stats_list = [] | |
| descriptor_funcs = { | |
| 'MolWt': Descriptors.MolWt, | |
| 'NumRotatableBonds': Descriptors.NumRotatableBonds, | |
| 'FormalCharge': lambda m: Chem.GetFormalCharge(m), | |
| } | |
| for dirpath, _, filenames in os.walk(root_path): | |
| if os.path.basename(dirpath) == 'raw_data': | |
| continue | |
| processed_file = None | |
| ids_col = None | |
| avg_components = np.nan | |
| std_components = np.nan | |
| # Look for a processed file to extract mixture info | |
| for fname in filenames: | |
| if fname.startswith("processed") and fname.endswith(file_extensions): | |
| processed_file = os.path.join(dirpath, fname) | |
| try: | |
| df_proc = pd.read_csv(processed_file) | |
| ids_cols = [col for col in df_proc.columns if re.search(r'ids', col, re.IGNORECASE)] | |
| if ids_cols: | |
| ids_col = ids_cols[0] | |
| comp_lengths = df_proc[ids_col].dropna().apply(ast.literal_eval).apply(len) | |
| avg_components = comp_lengths.mean() | |
| std_components = comp_lengths.std() | |
| except Exception as e: | |
| print(f"Error reading {processed_file}: {e}") | |
| break | |
| for fname in filenames: | |
| if not fname.endswith(file_extensions): | |
| continue | |
| fpath = os.path.join(dirpath, fname) | |
| dataset_name = os.path.normpath(fpath).split(os.sep)[-3] | |
| try: | |
| # CASE 1: IlThermo special handling | |
| if fname == "processed_IlThermoData.csv": | |
| compounds_path = os.path.join( | |
| root_path, "ionic-liquids", "processed_data", "compounds.csv" | |
| ) | |
| compounds_df = pd.read_csv(compounds_path) | |
| df = pd.read_csv(fpath) | |
| viscosity_ids = df[df['property'] == 'Viscosity']['cmp_ids'].apply(ast.literal_eval) | |
| conductivity_ids = df[df['property'] == 'Electrical conductivity']['cmp_ids'].apply(ast.literal_eval) | |
| subsets = { | |
| "Viscosity": { | |
| "data": compounds_df[compounds_df['compound_id'].isin( | |
| x for sublist in viscosity_ids for x in sublist | |
| )], | |
| "components": viscosity_ids | |
| }, | |
| "Electrical conductivity": { | |
| "data": compounds_df[compounds_df['compound_id'].isin( | |
| x for sublist in conductivity_ids for x in sublist | |
| )], | |
| "components": conductivity_ids | |
| } | |
| } | |
| for prop_name, sub_info in subsets.items(): | |
| smiles_list = sub_info["data"]['smiles'].dropna().astype(str).str.strip().unique() | |
| comp_lengths = sub_info["components"].apply(len) | |
| avg_comp = comp_lengths.mean() if not comp_lengths.empty else np.nan | |
| std_comp = comp_lengths.std() if not comp_lengths.empty else np.nan | |
| stats_list.extend( | |
| compute_molecule_stats_numeric(smiles_list, descriptor_funcs, "IlThermo_" + prop_name, avg_comp, std_comp) | |
| ) | |
| # CASE 2: Other datasets | |
| elif "ionic-liquids" not in dirpath: | |
| df = pd.read_csv(fpath) | |
| smiles_cols = [col for col in df.columns if re.search(r'smi|SMILES', col, re.IGNORECASE)] | |
| if not smiles_cols: | |
| continue | |
| smiles_set = set() | |
| for col in smiles_cols: | |
| valid_smiles = df[col].dropna().astype(str).str.strip() | |
| smiles_set.update(valid_smiles[valid_smiles != ''].unique()) | |
| stats_list.extend( | |
| compute_molecule_stats_numeric(smiles_set, descriptor_funcs, dataset_name, avg_components, std_components) | |
| ) | |
| except Exception as e: | |
| print(f"Skipped {fpath} due to error: {e}") | |
| return pd.DataFrame(stats_list) | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser(description="Extract SMILES and compute molecular statistics (with mixture info)") | |
| parser.add_argument("root_path", type=str, help="Root path to dataset directories") | |
| parser.add_argument("--output_dir", type=str, default="results", help="Directory to save results") | |
| args = parser.parse_args() | |
| os.makedirs(args.output_dir, exist_ok=True) | |
| print("Computing molecule statistics per dataset (with components info)...") | |
| stats_df = get_molecule_statistics_with_components(args.root_path) | |
| stats_df.to_csv(os.path.join(args.output_dir, "molecule_statistics.csv"), index=False) | |
| print(f"Saved molecule statistics to {args.output_dir}/molecule_statistics.csv") | |