| --- |
| license: apache-2.0 |
| library_name: catboost |
| tags: |
| - chemistry |
| - drug-discovery |
| - admet |
| - catboost |
| - chemprop |
| --- |
| |
| # Ark-ADMET: released CatBoost and Chemprop checkpoints for the TDC ADMET benchmark |
|
|
| This repository contains one five-seed CatBoost model set for each of the 22 Therapeutics Data Commons (TDC) ADMET endpoints, plus a separate five-seed Chemprop model set for `Lipophilicity_AstraZeneca`. |
|
|
| Models were evaluated using the official TDC splits and endpoint-specific metrics. CatBoost feature and hyperparameter selection used an internal scaffold split derived from the seed 1 training partition. The Chemprop model is an independent submission: its predictions are neither used as features by nor averaged with the CatBoost models. |
|
|
| ## Overview |
|
|
| - **22 CatBoost model sets**, one directory per TDC endpoint, using normalized lower-case directory names. |
| - **5 official TDC splits** per endpoint (`seed1` to `seed5`). Each `catboost.cbm` is fitted on the training partition from the corresponding split. |
| - **CatBoost prediction** is the mean over the five seeds. |
| - Classification: the score is the positive-class probability. |
| - Regression: the score is on the original label scale (the target scaler is inverted using `y_scaler.json`). |
| - **Optional feature blocks** (PaDEL 3D descriptors, GIN embeddings, MiniMol embeddings) are used only by the CatBoost endpoints that kept them; see the [endpoint table](#endpoints). |
| - **`match.json`** is an audit log from the CatBoost refit only: every saved test prediction matched the reported vector with a maximum absolute difference of 0 at `atol = 1e-5`. It is not needed to load or run a model. |
| - **Chemprop** lives under `chemprop/lipophilicity_astrazeneca/` and does not replace the CatBoost directory `lipophilicity_astrazeneca/`. |
|
|
| ## Quick start |
|
|
| Install the dependencies from `requirements.txt`: |
|
|
| | Package | Version | |
| | --- | --- | |
| | Python | 3.11.16 | |
| | `catboost` | 1.2.10 | |
| | `numpy` | 1.26.4 | |
| | `rdkit` | 2026.03.6 | |
| | `huggingface_hub` | 0.36.2 | |
|
|
| Download a snapshot and load an endpoint that uses only the fingerprint block: |
|
|
| ```python |
| from huggingface_hub import snapshot_download |
| from predict import EndpointModel |
| |
| repo_dir = snapshot_download("Eonark/Ark-ADMET") |
| model = EndpointModel.from_pretrained(repo_dir, "pgp_broccatelli") |
| print(model.predict(["CCO", "c1ccccc1"])) |
| ``` |
|
|
| Endpoints that kept extra blocks need those arrays supplied by the caller: |
|
|
| ```python |
| from predict import EndpointModel |
| |
| model = EndpointModel.from_pretrained(repo_dir, "herg") |
| |
| # padel: (n, 431) in feature_spec["padel_3d"]["columns"] order, |
| # or (n, n_3d) already sliced to the endpoint's kept columns |
| # gin: (n, 300) from the frozen GIN encoder described below |
| # minimol: (n, 512) from MiniMol v1; required when blend_weight > 0 |
| # or minimol_concat is true |
| pred = model.predict(smiles, padel=padel, gin=gin, minimol=minimol) |
| ``` |
|
|
| The MiniMol blend head is stored as `minimol_head.npz`, which is the preferred format and does not require scikit-learn. `minimol_head.joblib` is kept as the original fitted object. |
|
|
| ## Endpoints |
|
|
| The directory name is the second argument to `EndpointModel.from_pretrained`. The metric is the official TDC ADMET metric for that endpoint. Columns **3D**, **GIN**, **MiniMol concat**, and **Blend** show which optional feature blocks the endpoint uses (`3D` is the number of PaDEL 3D columns kept; `Blend` is `blend_weight`). |
|
|
| | Directory | TDC name | Task | Metric | 3D | GIN | MiniMol concat | Blend | |
| | --- | --- | --- | --- | --- | --- | --- | --- | |
| | `dili` | DILI | classification | ROC-AUC | 431 | | | 0.1 | |
| | `hia_hou` | HIA_Hou | classification | ROC-AUC | 0 | yes | | 0 | |
| | `bioavailability_ma` | Bioavailability_Ma | classification | ROC-AUC | 431 | yes | yes | 0.1 | |
| | `herg` | hERG | classification | ROC-AUC | 200 | yes | | 0.2 | |
| | `half_life_obach` | Half_Life_Obach | regression | Spearman | 431 | | | 0.8 | |
| | `cyp2d6_substrate_carbonmangels` | CYP2D6_Substrate_CarbonMangels | classification | PR-AUC | 100 | yes | | 0 | |
| | `cyp2c9_substrate_carbonmangels` | CYP2C9_Substrate_CarbonMangels | classification | PR-AUC | 100 | | | 0 | |
| | `cyp3a4_substrate_carbonmangels` | CYP3A4_Substrate_CarbonMangels | classification | ROC-AUC | 431 | | | 0.2 | |
| | `caco2_wang` | Caco2_Wang | regression | MAE | 431 | | | 0 | |
| | `clearance_microsome_az` | Clearance_Microsome_AZ | regression | Spearman | 431 | yes | yes | 0 | |
| | `vdss_lombardo` | VDss_Lombardo | regression | Spearman | 431 | | yes | 0 | |
| | `clearance_hepatocyte_az` | Clearance_Hepatocyte_AZ | regression | Spearman | 431 | | | 0 | |
| | `pgp_broccatelli` | Pgp_Broccatelli | classification | ROC-AUC | 0 | | | 0 | |
| | `bbb_martins` | BBB_Martins | classification | ROC-AUC | 431 | | | 0 | |
| | `ppbr_az` | PPBR_AZ | regression | MAE | 25 | yes | | 0.1 | |
| | `lipophilicity_astrazeneca` | Lipophilicity_AstraZeneca | regression | MAE | 431 | | | 0.3 | |
| | `ames` | AMES | classification | ROC-AUC | 0 | | yes | 0.2 | |
| | `ld50_zhu` | LD50_Zhu | regression | MAE | 0 | | | 0.2 | |
| | `solubility_aqsoldb` | Solubility_AqSolDB | regression | MAE | 50 | | | 0.2 | |
| | `cyp2c9_veith` | CYP2C9_Veith | classification | PR-AUC | 0 | yes | | 0.2 | |
| | `cyp3a4_veith` | CYP3A4_Veith | classification | PR-AUC | 200 | yes | | 0.1 | |
| | `cyp2d6_veith` | CYP2D6_Veith | classification | PR-AUC | 431 | yes | yes | 0 | |
| |
| ## Features |
| |
| ### Fingerprint block (all endpoints) |
| |
| Every model starts with the same **2563** fingerprint columns, in this order: |
| |
| 1. Morgan count, radius 2, 1024 bins (`GetHashedMorganFingerprint`) |
| 2. Avalon count, 1024 bins (`GetAvalonCountFP`) |
| 3. ErG, 315 |
| 4. 200 RDKit descriptors, named in `feature_spec.json` |
|
|
| `predict` computes this block directly from SMILES. Invalid SMILES are represented internally as missing features; callers should validate inputs and reject such rows rather than interpret the resulting predictions. |
|
|
| ### Optional blocks appended after the fingerprint |
|
|
| | Block | Used when | Width | Supplied via | |
| | --- | --- | --- | --- | |
| | PaDEL 3D | `n_3d > 0` | `n_3d` | `padel` | |
| | GIN | `gin` is true | 300 | `gin` | |
| | MiniMol | `minimol_concat` is true | 512 | `minimol` | |
|
|
| **PaDEL 3D (431 columns).** |
| - Only the 3D descriptor group is used (PaDEL flags `-3d -retain3d -retainorder`). |
| - Conformers come from RDKit ETKDGv3 with `randomSeed=1`, hydrogens kept, then MMFF optimization with `maxIters=200`. |
| - If the first embedding fails, it is retried with `useRandomCoords=True`. A molecule that still fails is omitted, and its 3D descriptors stay missing. |
| - Some endpoints keep a truncated, ranked subset of the 431 columns. The kept names and their indices into the 431-column list are stored in `feature_spec.json` under `endpoints.<name>.padel_columns` and `padel_index`. An endpoint with `n_3d` = 431 uses the full list in that order. |
|
|
| **GIN (300 dimensions).** |
| - Encoder: `dgllife.model.load_pretrained("gin_supervised_masking")`. |
| - Graphs: `mol_to_bigraph(add_self_loop=True, canonical_atom_order=False)` with `PretrainAtomFeaturizer` and `PretrainBondFeaturizer`. |
| - Readout: `dgl` `AvgPooling`. |
| - **The encoder weights are not bundled** in this repository. |
|
|
| **MiniMol (512 dimensions).** |
| - Embeddings come from public MiniMol v1. |
| - **The encoder weights are not bundled** in this repository. |
|
|
| ### MiniMol blend head |
|
|
| Where `blend_weight` is above 0, each seed also has a small head on the MiniMol vector: `StandardScaler`, one hidden layer of 128 ReLU units, then a linear output. The per-seed prediction is |
|
|
| ``` |
| (1 - alpha) * CatBoost + alpha * head |
| ``` |
|
|
| where `alpha` is the endpoint's `blend_weight`. If a row's MiniMol vector is not finite, that row keeps the CatBoost prediction unchanged. |
|
|
| ## Intended use and limitations |
|
|
| These checkpoints are intended for research, benchmark reproduction, and method comparison. They have not been validated for clinical, diagnostic, regulatory, or safety-critical use. Predictions may be unreliable for molecules and assay conditions outside the chemical and experimental domains represented by the TDC datasets. |
|
|
| The repository contains fitted parameters and inference helpers, not a complete ADMET decision system. Users remain responsible for input validation, applicability-domain assessment, uncertainty analysis, and experimental confirmation. |
|
|
| - **Unfeaturizable SMILES.** Two solubility-test molecules have no MiniMol vector and fall back to CatBoost: `CC1=CC=C[NH+2]([O-])[CH-]1` and `O=C(O)C1=C[NH+2]([O-])[CH-]C=C1`. |
| - **CYP2D6 substrate class weighting.** `CYP2D6_Substrate_CarbonMangels` was fitted with `scale_pos_weight` = 3. This weight is already baked into its trees; do not apply it again. |
| - **External encoders.** GIN and MiniMol embeddings must be computed by the caller with the specifications above. Mismatched encoder versions or featurization settings will change predictions. |
|
|
| ## Separate Chemprop submission (`Lipophilicity_AstraZeneca`) |
| |
| Directory: `chemprop/lipophilicity_astrazeneca/`. |
|
|
| This is the independent Chemprop 2.3.1 message-passing model reported alongside the CatBoost family. For each of the five official TDC splits, the checkpoint with minimum validation loss was retained; this is the only use of the official validation partition in this work. The released prediction is the mean of the five checkpoint outputs. Hyperparameters: message-passing depth 3, feed-forward depth 2, hidden dimension 300, dropout 0, up to 50 epochs, and molecular-graph input only (no extra RDKit descriptors). |
|
|
| | Item | Value | |
| | --- | --- | |
| | Task / metric | regression / MAE | |
| | Five-seed test MAE | 0.412 ± 0.012 (unrounded mean ± sample SD) | |
| | Relation to CatBoost | separate submission; not blended | |
| | CatBoost counterpart | still in `lipophilicity_astrazeneca/` | |
|
|
| The root `requirements.txt` covers CatBoost inference. Chemprop checkpoints require the packages in `chemprop/requirements.txt`. Load the five `seed{1-5}/best.pt` checkpoints with Chemprop 2.3.1 and average their predictions. See `chemprop/lipophilicity_astrazeneca/config.json` for the public hyperparameter record. |
|
|
| ## Repository layout |
|
|
| ``` |
| <endpoint>/config.json # CatBoost endpoints |
| <endpoint>/match.json # audit log only |
| <endpoint>/seed{1-5}/catboost.cbm |
| <endpoint>/seed{1-5}/y_scaler.json # regression only |
| <endpoint>/seed{1-5}/minimol_head.npz # blend_weight > 0 (preferred) |
| <endpoint>/seed{1-5}/minimol_head.joblib # same head, original object |
| chemprop/lipophilicity_astrazeneca/config.json |
| chemprop/lipophilicity_astrazeneca/seed{1-5}/best.pt |
| chemprop/lipophilicity_astrazeneca/seed{1-5}/metrics.json |
| chemprop/requirements.txt |
| feature_spec.json |
| versions.json |
| predict.py # CatBoost inference helper |
| requirements.txt |
| ``` |
|
|
| ## License |
|
|
| Released under the Apache License 2.0. Copyright 2026 ADMET endpoint models contributors. |
|
|
| Third-party fingerprints, PaDEL, GIN, and MiniMol remain under their own licenses. Their pretrained weights are not included in this repository. |
|
|
