Skills › Scientific Agent Skills › deepchem › predict_solubility.py
predict_solubility.py
221 lines · skills/deepchem/scripts/predict_solubility.py at 5f2bc03
1#!/usr/bin/env python3
2"""
3Molecular Solubility Prediction Script
5This script trains a model to predict aqueous solubility from SMILES strings
6using the Delaney (ESOL) dataset as an example. Can be adapted for custom datasets.
8Usage:
9 python predict_solubility.py --data custom_data.csv --smiles-col smiles --target-col solubility
10 python predict_solubility.py # Uses Delaney dataset by default
11"""
13import argparse
14import deepchem as dc
15import numpy as np
16import sys
17from _common import load_csv, validate_splits, evaluate_splits
20def train_solubility_model(data_path=None, smiles_col='smiles', target_col='measured log solubility in mols per litre', n_epochs=50):
21 """
22 Train a solubility prediction model.
24 Args:
25 data_path: Path to CSV file with SMILES and solubility data. If None, uses Delaney dataset.
26 smiles_col: Name of column containing SMILES strings
27 target_col: Name of column containing solubility values
29 Returns:
30 Trained model, test dataset, and transformers
31 """
32 print("=" * 60)
33 print("DeepChem Solubility Prediction")
34 print("=" * 60)
36 # Load data
37 if data_path is None:
38 print("\nUsing Delaney (ESOL) benchmark dataset...")
39 tasks, datasets, transformers = dc.molnet.load_delaney(
40 featurizer=dc.feat.CircularFingerprint(radius=2, size=2048),
41 splitter='scaffold'
42 )
43 train, valid, test = datasets
44 else:
45 print(f"\nLoading custom data from {data_path}...")
46 featurizer = dc.feat.CircularFingerprint(radius=2, size=2048)
47 dataset = load_csv(data_path, [target_col], smiles_col, featurizer)
48 if not np.all(dataset.w > 0):
49 raise ValueError('Solubility training requires a measured target for every row')
51 # Split data
52 print("Splitting data with scaffold splitter...")
53 splitter = dc.splits.ScaffoldSplitter()
54 train, valid, test = splitter.train_valid_test_split(
55 dataset,
56 frac_train=0.8,
57 frac_valid=0.1,
58 frac_test=0.1
59 )
61 # Normalize data
62 print("Normalizing targets using training statistics...")
63 transformers = [
64 dc.trans.NormalizationTransformer(
65 transform_y=True,
66 dataset=train
67 )
68 ]
69 for transformer in transformers:
70 train = transformer.transform(train)
71 valid = transformer.transform(valid)
72 test = transformer.transform(test)
74 tasks = [target_col]
76 print(f"\nDataset sizes:")
77 print(f" Training: {len(train)} molecules")
78 print(f" Validation: {len(valid)} molecules")
79 print(f" Test: {len(test)} molecules")
81 validate_splits((train, valid, test), "regression", len(tasks))
83 # Create model
84 print("\nCreating multitask regressor...")
85 model = dc.models.MultitaskRegressor(
86 n_tasks=len(tasks),
87 n_features=2048, # ECFP fingerprint size
88 layer_sizes=[128, 64],
89 dropouts=0.25,
90 learning_rate=0.001,
91 device="cpu",
92 batch_size=50
93 )
95 # Train model
96 print("\nTraining model...")
97 model.fit(train, nb_epoch=n_epochs)
98 print("Training complete!")
100 # Evaluate model
101 print("\n" + "=" * 60)
102 print("Model Evaluation")
103 print("=" * 60)
105 evaluate_splits(model, (train, valid, test), "regression", transformers)
107 return model, test, transformers
110def predict_new_molecules(model, smiles_list, transformers=None):
111 """
112 Predict solubility for new molecules.
114 Args:
115 model: Trained DeepChem model
116 smiles_list: List of SMILES strings
117 transformers: List of data transformers to apply
119 Returns:
120 Array of predictions
121 """
122 print("\n" + "=" * 60)
123 print("Predicting New Molecules")
124 print("=" * 60)
126 # Featurize new molecules
127 featurizer = dc.feat.CircularFingerprint(radius=2, size=2048)
128 features = featurizer.featurize(smiles_list)
130 if any(np.asarray(feature).shape != (2048,) for feature in features):
131 raise ValueError('Cannot predict invalid or unsupported SMILES')
133 # Create dataset
134 new_dataset = dc.data.NumpyDataset(X=features)
136 # Both training paths normalize with transform_y=True, so the model learned
137 # in z-scored target space. Hand the transformers to predict() so it
138 # untransforms the output back to target units -- transforming the dataset
139 # instead would do nothing (these transformers touch y, not X, and a
140 # prediction dataset has no y) and leave the printed numbers z-scored.
141 predictions = np.asarray(model.predict(new_dataset, transformers=transformers or [])).reshape(len(smiles_list), 1)
143 # Display results
144 print("\nPredictions:")
145 for smiles, pred in zip(smiles_list, predictions):
146 print(f" {smiles:30s} -> {pred[0]:.3f} target units")
148 return predictions
151def main():
152 parser = argparse.ArgumentParser(
153 description='Train a molecular solubility prediction model'
154 )
155 parser.add_argument(
156 '--data',
157 type=str,
158 default=None,
159 help='Path to CSV file with molecular data'
160 )
161 parser.add_argument(
162 '--smiles-col',
163 type=str,
164 default='smiles',
165 help='Name of column containing SMILES strings'
166 )
167 parser.add_argument(
168 '--target-col',
169 type=str,
170 default='solubility',
171 help='Name of column containing target values'
172 )
173 parser.add_argument(
174 '--predict',
175 nargs='+',
176 default=None,
177 help='SMILES strings to predict after training'
178 )
180 parser.add_argument('--epochs', type=int, default=50)
181 args = parser.parse_args()
182 if args.epochs < 1:
183 parser.error('--epochs must be positive')
185 # Train model
186 try:
187 model, test_set, transformers = train_solubility_model(
188 data_path=args.data,
189 smiles_col=args.smiles_col,
190 target_col=args.target_col,
191 n_epochs=args.epochs
192 )
193 except Exception as e:
194 print(f"\nError during training: {e}", file=sys.stderr)
195 return 1
197 # Make predictions on new molecules if provided
198 if args.predict:
199 try:
200 predict_new_molecules(model, args.predict, transformers)
201 except Exception as e:
202 print(f"\nError during prediction: {e}", file=sys.stderr)
203 return 1
204 else:
205 # Example predictions
206 example_smiles = [
207 'CCO', # Ethanol
208 'CC(=O)O', # Acetic acid
209 'c1ccccc1', # Benzene
210 'CN1C=NC2=C1C(=O)N(C(=O)N2C)C', # Caffeine
211 ]
212 predict_new_molecules(model, example_smiles, transformers)
214 print("\n" + "=" * 60)
215 print("Complete!")
216 print("=" * 60)
217 return 0
220if __name__ == '__main__':
221 sys.exit(main())