Molecular graph learning, chemical descriptors, and calibrated uncertainty for aqueous solubility.
Mo Jaffri · Chemical Engineering / Scientific Machine Learning
An auditable, CPU-runnable study of a practical molecular property: log10 aqueous solubility in mol/L. It compares a bond-aware PyTorch message-passing network with RDKit baselines on the public ESOL dataset. The central question is whether learned molecular representations add value over chemical descriptors when data are scarce and test molecules have different scaffolds.
The contribution is the complete experiment: structure cleaning, reproducible partitions, strong baselines, three controlled ablations, independent uncertainty calibration, and inspectable predictions. This is a chemistry-informed modeling study, with optional geometric features; it does not impose a thermodynamic equation or simulate molecular dynamics.
The full measured results, per-seed metrics, and paired ablations are in the benchmark report. Every reported value is generated from saved predictions. No literature score is presented as a directly comparable result.
Measured on Windows CPU, three seeds per split. RMSE is in log10(mol/L); ± denotes sample SD across the three runs, not a confidence interval.
| Model | Random RMSE ↓ | Scaffold RMSE ↓ |
|---|---|---|
| Training mean | 2.003 ± 0.113 | 2.915 ± 0.033 |
| Descriptor ridge | 0.979 ± 0.028 | 1.076 ± 0.062 |
| Descriptor RF | 0.705 ± 0.064 | 0.825 ± 0.024 |
| Morgan RF | 1.199 ± 0.050 | 2.147 ± 0.042 |
| Graph-only GNN | 0.644 ± 0.064 | 0.785 ± 0.059 |
| GNN + descriptors | 0.593 ± 0.039 | 0.764 ± 0.063 |
| GNN + descriptors, bond features masked | 0.598 ± 0.033 | 0.762 ± 0.100 |
| GNN + descriptors + 3D | 0.605 ± 0.084 | 0.807 ± 0.091 |
GNN + descriptors reduced mean random-split RMSE by 15.9% relative to the descriptor forest in this fixed experiment. Its scaffold RMSE was 0.764 ± 0.063. These are within-project comparisons on the documented partitions.
At nominal 90% coverage, its absolute conformal intervals covered 87.5% of random-split test targets and 91.4% of scaffold-split targets on average, with widths 1.967 and 2.619 logS respectively. Coverage varied across seeds; the method does not promise 90% coverage under scaffold shift.
The 48 configurations took 17.4 minutes of summed model-training time on this CPU, excluding featurization and other overhead. The three-epoch smoke check is only a functionality test.
Verification: 26 tests passed locally; package installation/build and CLI train/evaluate/predict/download were exercised. The full-run audit independently verified all 48 configurations. CPU checkpoint reload reproduced predictions exactly. The GitHub workflow additionally audits the published results and trains a small GNN on Linux. See findings and interpretation.
Additional evidence: a Tanimoto kernel-ridge follow-up adds six validation-selected comparisons. A CPU batching study measures 19.5× greater throughput for batch 128 versus batch 1 on one preloaded GNN member, with prediction differences below 10⁻⁶. These supplement the original 48-run study; all scope and timing exclusions are documented.
| Component | Implementation |
|---|---|
| Public data | ESOL / Delaney, 1,128 source rows; checksum-verified download |
| Cleaning | RDKit canonical isomeric SMILES; duplicate target averaging before splitting; rejection and duplicate audit |
| Baselines | Training mean; standardized descriptor ridge; descriptor random forest; Morgan fingerprint random forest |
| Graph network | Pure PyTorch bond-conditioned message passing, residual updates, sum + mean molecular pooling |
| Chemistry features | 12 descriptors reflecting molecular size, polarity, hydrogen bonding, flexibility and aromaticity |
| Geometry ablation | ETKDGv3 + MMFF94; six rotation/translation-invariant shape descriptors and a missingness flag |
| Uncertainty | Three independently initialized GNNs; separate split-conformal calibration, absolute and spread-scaled intervals |
| Evaluation | Random and scaffold partitions; three seeds; RMSE, MAE, R², coverage, width, interval score |
| Ablations | Remove global descriptors; mask bond features; add conformer shape features |
| Reproducibility | Pinned dependencies, JSON configs, dataset/source hashes, split IDs, learning curves, CLI, tests, GitHub Actions |
Python 3.11–3.13 is supported; the reference run uses 3.12. Commands below run from this repository's root.
python -m venv .venvActivate on macOS/Linux with source .venv/bin/activate; on Windows PowerShell use .venv\Scripts\Activate.ps1. If script activation is disabled, invoke .venv\Scripts\python.exe directly.
python -m pip install --upgrade pip
python -m pip install torch==2.7.1 --index-url https://download.pytorch.org/whl/cpu
python -m pip install -e ".[dev]"
solubility-bench download
pytest -qOn macOS, install torch==2.7.1 from the default PyPI index instead of the CPU wheel index above. Reference experiments were measured on Windows CPU. The status badge links to the latest Linux tests, data audit and smoke-training checks.
Short functional check (three epochs, intentionally not a performance benchmark):
solubility-bench benchmark --config configs/smoke.json --artifacts artifacts/smoke --reports reports/smokeFull reference experiment (48 model/split/seed configurations; 72 GNN members):
solubility-bench benchmark --config configs/benchmark.json --artifacts artifacts/reproduction --reports reports/reproductionThe full run includes conformer generation, 24 classical fits/selections, and 24 three-member GNN ensembles. Expect substantially longer than the smoke check on a laptop; measured training times are in the result table. Feature generation is cached using the dataset hash, featurization code hash, RDKit version, and conformer seed. The default benchmark report is bundled for immediate inspection; the reproduction command writes a fresh report separately.
The top-level requirements file pins the scientific packages. The Windows CPU environment snapshot also records transitive and development dependencies from the measured run. To reproduce that environment, install the CPU torch wheel first, then that snapshot, then pip install -e . --no-deps. Use the main requirements for other operating systems.
For an interrupted run, repeat its command with --resume. Existing checkpoints are reused only when source, config, and device match. Use a new output directory when changing any of these. Model checkpoints and feature caches are ignored by Git.
The source release archive includes the code, tests, documentation and measured reports. It excludes the Python environment, downloaded source CSV, caches and trained checkpoints. Generate checkpoints with the commands below.
solubility-bench train --config configs/benchmark.json --model gnn_descriptors --split scaffold --seed 17 --artifacts artifacts/example --reports reports/example
solubility-bench evaluate --checkpoint artifacts/example/model.joblib --reports reports/example
solubility-bench predict --checkpoint artifacts/example/model.joblib --smiles "CCO" "CC(=O)Oc1ccccc1C(=O)O" --output predictions.csvtrain uses training labels, validation labels for model selection, and calibration labels for intervals. evaluate scores the held-out test partition. Predictions are in log10(mol/L); interval bounds use the configured nominal coverage. They are research outputs, not experimentally verified values for new compounds.
Equivalent entry points: python -m solubility_bench.cli, python scripts/train.py, python scripts/evaluate.py, and python scripts/benchmark.py. Run solubility-bench --help for options. Load only trusted local joblib checkpoints/caches: the format can execute Python on loading.
The same model and graph batching code supports --device cuda. Install a CUDA-compatible PyTorch 2.7.1 wheel appropriate for your driver using the official version-specific instructions, then install the project. The CPU reference is tested; CUDA execution is supported by the code but was not validated on GPU hardware. Seeds and deterministic-mode warnings are enabled, but identical floating-point results across platforms/devices are not promised.
- Verify and clean. Read only SMILES and measured logS. Ignore the source's
ESOL predicted log solubilityand all precomputed descriptors. Merge canonical duplicate structures by averaging their labels before splitting. Keep source-row provenance and duplicate label ranges. Do not silently normalize salts, tautomers, or protonation states. - Partition without labels. Target 65% training / 10% validation / 10% calibration / 15% test. Three fixed seeds: 17, 42, 73. Random splits shuffle unique structures. Scaffold splits keep each nonchiral Bemis–Murcko scaffold intact. Largest groups are assigned first to the largest remaining count deficit; seed breaks equal-size ties. All acyclic molecules share the empty scaffold. Large groups can move realized fractions away from targets. This is a custom four-way split, not DeepChem's standard split.
- Fit on training only. Descriptor/shape imputation, descriptor scaling, and GNN target normalization use training observations only. No refit on validation or calibration data occurs after selection.
- Select on validation only. Ridge alpha ∈ {0.1, 1, 10, 100}; forest minimum leaf size ∈ {1, 3, 5}, 250 trees. Each GNN uses fixed hyperparameters and validation early stopping independently within its three-member ensemble. All runs share partitions within each split/seed.
- Calibrate once the model is frozen. Use the dedicated calibration partition for 90% prediction intervals. Ensemble/forest spread is a diagnostic, not a calibrated uncertainty estimate by itself.
- Evaluate the fixed matrix. Compare all eight model variants on held-out test molecules. Report every seed, mean ± sample SD, and paired ablation differences. The SD is descriptive variation, not a statistical confidence interval; repeated test sets overlap.
The methodology note gives architecture, equations, calibration details, and limitations. The data card records provenance and reuse considerations.
Solubility reflects the balance between a molecule's interactions with water and the cost of leaving its original phase. Molecular size, lipophilicity, hydrogen bonding, and solid-state packing all matter. That makes chemical descriptors useful inductive information, particularly on a small dataset.
For ideal dilute dissolution at equilibrium, with a defined standard state and a pure solid, ΔG°diss ≈ −RT ln(S/c°). This motivates a logarithmic concentration target, but this dataset does not provide the phase, pH, temperature and activity information needed to turn every label into a reliable free energy. The models therefore predict empirical logS directly. They do not claim a physically constrained loss, solvation free-energy calculation, or a full description of crystal packing.
The graph network learns local chemical environments; global descriptors supply chemically interpretable information; the geometry experiment tests whether a single approximate molecular shape adds signal. The ablations let the evidence decide whether these additions help.
- Results and comparison tables
- All per-run metrics and aggregate metrics
- Paired ablations
- Data cleaning and conformer audit
- Per-run
split.json,training.json,metrics.json, andpredictions.csvinreports/benchmark/<split>/seed-<seed>/<model>/ - Calibration plot and scaffold parity/similarity diagnostics
The parity and similarity plots always use GNN + descriptors and the first seed. The plotted example is fixed, not chosen for attractive performance. Calibration curves are diagnostics; the primary 90% setting remains fixed.
Independently verify the bundled report against source labels and prediction CSVs:
python scripts/audit_results.pyAfter producing a new benchmark, export its calibration predictions and audit it:
python scripts/audit_results.py --reports reports/reproduction --artifacts artifacts/reproductionThe auditor reconstructs partitions and checks source labels, calibration quantiles, interval bounds, test counts and metrics. The reference pass is recorded in audit.json.
configs/ Full experiment and short smoke configuration
src/solubility_bench/
data.py Source verification, cleaning, scaffold splits
features.py RDKit descriptors, Morgan, graphs, conformers
models.py Bond-aware PyTorch message passing
training.py Baselines, ensembles, early stopping, inference
uncertainty.py Finite-sample conformal intervals and metrics
experiment.py Auditable artifacts and evaluation boundary
reporting.py Tables, ablations, calibration and diagnostic plots
cli.py Download/train/evaluate/predict/benchmark/report
tests/ Invariance, leakage, calibration, checkpoint tests
scripts/ Thin training/evaluation/benchmark entry points
reports/benchmark/ Measured results, predictions, partitions and plots
docs/ Methodology, data card, resume and interview notes
data/raw/ Downloaded source (ignored by Git)
artifacts/ Local trained models (ignored by Git)
- ESOL is small and chemically uneven. Four-way splitting reduces model-fitting data further. A descriptor model can outperform a graph model; that is a useful result, not a reason to discard the baseline.
- Scaffold separation limits one form of leakage but does not enforce low chemical similarity or reproduce a prospective medicinal-chemistry campaign. Common scaffolds and the large acyclic group can dominate training.
- Standard split-conformal guarantees require exchangeability. They do not automatically hold under scaffold distribution shift; observed coverage and interval width must be read together. No conditional guarantee is made for a scaffold or compound.
- Duplicate averaging changes the effective task and can hide experimental variation. No assay noise, pH, temperature, polymorph, or ionization context is explicitly modeled.
- A single optimized conformer is neither an aqueous ensemble nor a crystal structure. Missing/nonconverged conformers are imputed with training medians and explicitly flagged; their labels are retained.
- The GNN does not encode stereochemistry explicitly. It uses atom/bond categories and graph structure, not 3D equivariant message passing. Shape descriptors are invariant summaries, not learned geometric interactions.
- Hyperparameter search budgets differ between classical models and GNNs. This is a fixed-budget portfolio benchmark, not a claim of optimal architectures or state-of-the-art performance.
Useful extensions: scaffold-group bootstrap intervals; more independent scaffold partitions; an external solubility dataset with an overlap audit; broader validation-only GNN tuning; multi-conformer features; and a separate QM9 energy/force task with explicit symmetry and physical targets. See the interview and extension guide.
For molecular ML or computational drug-discovery applications, place Solubility Bench above DistillTwin and Titrate within Projects, while keeping substantive research in Research Experience. For process engineering, keep DistillTwin first. If space permits only two projects for a scientific ML application, prefer Solubility Bench + DistillTwin unless Titrate has stronger demonstrated use or technical depth. This ordering is based on the described roles of those projects, not an audit of their code.
Resume bullets and a guide to explaining the work are in RESUME_AND_INTERVIEW.md. Use measured claims and be ready to explain or modify the implementation.
- Delaney, J. S. (2004). ESOL: Estimating Aqueous Solubility Directly from Molecular Structure.
- Wu et al. (2018). MoleculeNet: a benchmark for molecular machine learning.
- DeepChem's ESOL loader and source definition.
- Gilmer et al. (2017). Neural Message Passing for Quantum Chemistry. The network here is a small original implementation in this framework, not a reproduction of their architecture/results.
- RDKit documentation: fingerprints and ETKDG.
- Angelopoulos & Bates (2021). A Gentle Introduction to Conformal Prediction and Distribution-Free Uncertainty Quantification.
Project code: MIT. Dataset attribution and terms remain with the original sources; the dataset is downloaded separately and is not relicensed by this project. This independent portfolio project has no affiliation with D. E. Shaw Research or any employer.