CoTAR — bond-order predictor (evaluation release)

CoTAR predicts chemical bond types (single / double / triple / aromatic) directly from the 3D coordinates of a molecule using a graph neural network.

This repository provides a reduced release for academic evaluation. It includes the core inference implementation, a pinned software environment, demonstration structures, a minimal training example, and runnable examples. Machine-readable numerical data underlying Figures 2 and 3 and Tables 1–3 of the associated article are provided with the article as Supporting Information. See LICENSE for usage terms.

Installation

Requires Python 3.11–3.13.

With uv (recommended)

A pinned environment is provided via uv.lock. From the project root:

uv sync                                  # create .venv and install exact pinned deps
uv run python examples/example_inference.py

uv run <cmd> executes inside the project environment. The lockfile pins a CUDA 12.8 PyTorch build (torch < 2.11, which still supports NVIDIA V100 / sm_70); it also runs on CPU when no GPU is present. The example auto-selects CUDA when available.

With pip

pip install .

This installs the dependencies used by inference and the minimal training example: torch, numpy, vesin, and rdkit.

Quick start

from rdkit import Chem
from cotar import COTAR

# Load the bundled pretrained model.
model = COTAR.from_pretrained()        # use device="cuda" if available

# A molecule with 3D coordinates (bond orders are ignored / re-predicted).
mol = Chem.MolFromMolFile("data/demo/ethanol.sdf", removeHs=False)

# Predict bonds.
result = model.predict(mol)
for i, j, bond_type in result["bonds"]:
    print(f"atom {i} - atom {j}: {bond_type}")

# Or get a new RDKit molecule with the predicted bonds applied.
new_mol = model.update_bonds(mol)
print(Chem.MolToSmiles(new_mol))

predict() returns a dictionary with:

• bonds: list of (atom_i, atom_j, bond_type_str) for predicted bonds
• edge_index: (N_edges, 2) atom-pair indices in the neighbor list
• pred_labels: (N_edges,) predicted bond-type indices
• pred_probs: (N_edges, 5) per-edge probabilities [SINGLE, DOUBLE, TRIPLE, AROMATIC, NoBond]

You can also pass a dictionary with symbols and positions (and optionally cell/pbc for periodic systems) instead of an RDKit Mol.

Demonstration data

data/demo/ holds a few small molecules in RDKit-readable structure files (.sdf, with 3D coordinates and reference bonds). They illustrate the input format expected by the model and are used by the inference and minimal training examples. They are a tiny demonstration sample, not the training dataset used for the paper.

Supported elements. This reduced release is restricted to the elements in the demonstration data — H, C, and O. The bundled model and the element tables have been pruned accordingly; other elements are not supported.

Bundled model. The bundled default model is not the model from the paper. It is a small model trained for only 10 epochs on the demonstration molecules, provided solely to exercise the code end to end — its predictions are not expected to be accurate.

Inference example

examples/example_inference.py is a minimal single-molecule walkthrough you can run directly:

python examples/example_inference.py

Minimal training example

examples/example_training.py trains a small model from scratch using all four bundled SDF files, saves a checkpoint with its configuration embedded, reloads it through the standard inference API, and runs one prediction:

uv run python examples/example_training.py

The default 10-epoch run writes an ignored artifact to outputs/demo_training/demo_model.pt. It can be loaded directly:

from cotar import COTAR

model = COTAR.from_checkpoint("outputs/demo_training/demo_model.pt")

This example trains only the bond-classification path: atom-property losses and the van der Waals prior are disabled because the bundled sample only supplies the bond labels needed for this demonstration. It uses every structure for training and does not create a validation split. The four single-conformer structures are far too small for model assessment and contain no triple bonds; the resulting checkpoint is only evidence that data preparation, optimization, checkpointing, and inference run end to end. It must not be interpreted as a reproduction of the paper model, evaluation, or benchmark.