mlxmolkit — GPU-accelerated molecular toolkit on Apple Silicon

Port of nvMolKit (CUDA) to Apple Metal via MLX. Three pipelines:

Molecular Clustering — Morgan FP → Tanimoto similarity → Butina clustering
3D Conformer Generation — DG (4D) → ETK (3D) → MMFF94 optimization
PM6_D semi-empirical SCF — full d-orbital NDDO (S/P/Cl/Br/I) with PM6-D3H4 corrections

What's new

Semi-empirical SCF on Apple Silicon — 7 methods (RM1, AM1, PM3, PM6, PM6_SP, PM6_D, AM1*) plus PM6-D3H4 post-SCF corrections — bit-exact to PYSEQM for PM6_D, with no PYSEQM/PyTorch runtime dependency. Every entry point is covered by tests/test_{methods_api,pm6_d_native,pm6_d3h4,pyseqm_port,rm1_scf}.py (83 tests total).

Method	Coverage	HoF (H2O, kcal/mol)	Status
RM1	H, C, N, O, F, P, S, Cl, Br, I	-57.81	tested
AM1	H, C, N, O	-59.22	tested
PM3	H, C, N, O, F, P, S, Cl, Br, I	-53.19	tested
PM6 / PM6_SP	H, C, N, O, F, P, S, Cl, Br, I (sp-only)	-54.19	tested
PM6_D	+ d-orbitals on P, S, Cl, Br	bit-exact vs PYSEQM	tested
AM1* / RM1*	H, C, N, O (*-variants)	-53.71 / -54.47	tested
PM6-D3H4	D3 dispersion + H4 H-bond + HH repulsion	post-SCF correction	tested

PM6_D SCF in pure NumPy — per-pair W tensor matches PYSEQM to 2.66e-15 (machine epsilon). 27/27 SCF charge tests pass against frozen PYSEQM/MOPAC references.
Full d-orbital support — P, S, Cl (qn=3) and Br (qn=4) with the 22-integral local frame, rotated to molecular frame via Wigner D-matrices. Covers YH, YX, YY pair types.
PM6-D3H4 post-SCF corrections — Grimme D3 dispersion + Rezáč–Hobza H4 hydrogen-bond + HH-repulsion.
DIIS + adaptive damping SCF — converges reliably on hard cases (CCl4, SF6, DMSO) where plain mixing freezes the wrong basin.
numba JIT + einsum hot kernels — w_withquaternion, GenerateRotationMatrix, Rotate2Center2Electron accelerated; ~2× end-to-end per pair (3.3 → 1.6 ms per S-C pair on M3 Max).
Bit-exactness regression suite — 23 frozen-reference tests guard the numerical invariants of the vendored PYSEQM port.

Public API

from mlxmolkit.nddo import (
    # SCF
    nddo_energy, nddo_energy_batch,
    # PM6-D3H4 corrections
    pm6_d3h4_correction, d3_energy, h4_energy, hh_repulsion,
    # Parameters
    METHOD_PARAMS, get_params, ElementParams,
)

# Bit-exact primitives (vendored from PYSEQM)
from mlxmolkit.nddo._pyseqm_port import (
    diatom_overlap_matrixD, two_elec_two_center_int,
    qn_int, qnD_int,
)

# Example: PM6_D total energy with D3H4 corrections
result = nddo_energy(atoms=[16, 1, 1], coords=[[0,0,0], [0.97,0,0.93], [-0.97,0,0.93]], method='PM6_D')
correction = pm6_d3h4_correction(atoms=[16, 1, 1], coords=...)

Installation

pip install mlxmolkit-rdkit

Requires macOS with Apple Silicon (M1/M2/M3/M4). RDKit is needed for molecular input:

conda install -c conda-forge rdkit
pip install mlxmolkit-rdkit

Quick Start

from mlxmolkit import generate_conformers_nk, butina_tanimoto_mlx

# 3D conformers: 100 molecules x 10 conformers each
result = generate_conformers_nk(
    smiles_list=["c1ccccc1", "CC(=O)O", "CC(=O)Oc1ccccc1C(=O)O"],
    n_confs_per_mol=10,
    run_mmff=True,
)

# Clustering: 150k+ molecules
import mlx.core as mx
result = butina_tanimoto_mlx(mx.array(fp_bytes), cutoff=0.4)

Features

Conformer Generation — Drop-in replacement for RDKit's ETKDG (EmbedMolecules). Supports ETKDG, ETKDGv2, ETKDGv3, srETKDGv3, KDG, ETDG, ETDGv2, and pure DG.
MMFF94 Optimization — GPU-accelerated force field optimization (MMFFOptimizeMoleculesConfs). All 7 MMFF energy terms with fused Metal kernel. Full BFGS or L-BFGS in-kernel (zero CPU round-trips).
Molecular Clustering — Butina clustering at 150k+ molecules with divide-and-conquer memory management.
N x k Parallel — Generate k conformers for N molecules simultaneously. Constraints shared across conformers (conf_to_mol indirection, 50% memory savings).

Performance

Conformer Generation (N=20 molecules, k=50 conformers = 1000 total)

Pipeline	Time	Throughput	GPU Memory
DG only	1.00s	1,002 conf/s	2.6 MB
DG + ETK	1.11s	900 conf/s	2.6 MB
DG + ETK + MMFF	1.63s	614 conf/s	5.1 MB

Conformer Memory Scaling (DG + ETK + MMFF, batch=500)

Conformers	Batch	GPU (BFGS)	GPU (L-BFGS)	Time	Throughput
1,000	1000	5.1 MB	2.9 MB	0.43s	2,342/s
2,000	500	2.6 MB	1.5 MB	1.43s	1,402/s
4,000	500	2.6 MB	1.5 MB	1.91s	2,094/s
10,000	500	2.6 MB	1.5 MB	4.82s	2,075/s

GPU memory stays constant regardless of total conformers thanks to divide-and-conquer batching.

Scale Tests

Scale	Pipeline	Time	Throughput	Convergence
N=1000, k=10	DG + ETK	5.0s	2,017 conf/s	99.7%
N=1000, k=10	DG + ETK + MMFF	8.0s	1,250 conf/s	99.7%
N=10000, k=1	DG + ETK	17.7s	565 conf/s	99.6%
N=10000, k=1	DG + ETK + MMFF	37.9s	264 conf/s	99.6%
N=10000, k=10	DG + ETK	38.1s	2,625 conf/s	99.7%
N=10000, k=10	DG + ETK + MMFF	67.9s	1,473 conf/s	99.7%

100,000 conformers in a single GPU batch. All stages on Metal (including MMFF94 — zero RDKit post-processing).

Batch Size Impact (N=20, k=50, C=1000)

Batch	Batches	Time	conf/s
100	10	0.62s	1,610
500	2	0.29s	3,394
1000+	1	0.22s	4,442

Larger batches = fewer kernel launches = higher throughput. Auto-sizing (default) picks the largest batch that fits in free memory.

GPU Memory per Conformer

Atoms	DG (4D)	ETK (3D)	MMFF (BFGS)	MMFF (L-BFGS)
5	1.9 KB	1.4 KB	1.3 KB	1.4 KB
12	4.4 KB	3.3 KB	6.0 KB	3.3 KB
21	7.6 KB	5.7 KB	17.2 KB	5.7 KB
30	10.8 KB	8.1 KB	34.1 KB	8.1 KB
50	18.0 KB	13.5 KB	92.0 KB	13.5 KB
64	23.1 KB	17.3 KB	149.2 KB	17.3 KB

MMFF BFGS memory grows as O(n^2) due to the dense Hessian (n_atoms x 3)^2. BFGS is faster than L-BFGS at all typical drug-like sizes (up to 74 atoms with H) because the better curvature information requires fewer iterations. L-BFGS is only needed for very large molecules (>150 atoms) where the Hessian exceeds ~1 MB per conformer.

Molecule	Atoms (with H)	BFGS	L-BFGS	Winner
Methane	5	0.255s	0.215s	BFGS
Benzene	12	0.213s	0.222s	~tie
Aspirin	21	0.241s	0.230s	~tie
Testosterone	49	0.364s	0.335s	BFGS
Cholesterol	74	0.590s	0.486s	BFGS

Recommendation: use BFGS (default) for all molecules <150 atoms with H. The pipeline auto-switches to L-BFGS at 150+ atoms (mmff_use_lbfgs=None, the default).

Important: Always add explicit hydrogens (Chem.AddHs) before conformer generation. Convergence is significantly better with explicit H because the distance geometry constraints are more complete and the force field terms (bond/angle/torsion) are fully defined. The pipeline calls AddHs automatically.

With 64 GB unified memory, a single batch can hold:

Molecule size	DG/ETK	MMFF (BFGS)	MMFF (L-BFGS)
12 atoms	~9.8M conformers	~7.4M	~9.8M
30 atoms	~3.9M conformers	~1.3M	~3.9M
64 atoms	~1.8M conformers	~300K	~1.8M

The divide-and-conquer queue automatically splits into multiple batches when total exceeds free memory.

Clustering (Enamine REAL subset, Apple M3 Max)

N	Fused sim→CSR	Butina	Total	vs RDKit	Memory
20k	0.26s	0.09s	0.35s	152x	0.1 MB
50k	1.26s	0.36s	1.62s	—	0.5 MB
100k	4.87s	0.97s	5.84s	—	1.3 MB
150k+	blockwise	—	scales	—	bounded

ETKDG Variant Comparison (N=20, k=50)

Variant	conf/s	Convergence
DG	1,005	100.0%
KDG	910	99.1%
ETDG	1,009	100.0%
ETDGv2	1,008	100.0%
ETKDG	907	99.1%
ETKDGv2	904	99.1%
ETKDGv3	895	99.1%
srETKDGv3	914	100.0%

RDKit vs mlxmolkit conformer quality (N=20, k=10)

After the ETK/MMFF packing fixes and batched MLX Horn alignment, the current RDKit-vs-mlxmolkit comparison on the CHEESE charge-training subset reports:

Setting	Value
RDKit reference	ETKDGv3 + MMFF94
mlxmolkit path	ETKDGv3 + MMFF94
Mean best shape Carbo	0.934633
Mean best electrostatic Carbo	0.996936
Mean best combined score	0.966551
Median best heavy-atom RMSD	0.449579 A

The scoring path aligns all mlxmolkit conformers to all RDKit conformers with batched MLX Horn quaternion alignment, then evaluates paired shape Carbo and electrostatic Carbo in one MLX tensor pass per molecule. It avoids per-pair CPU Kabsch/SVD loops and repeated tiny CHEESE kernel launches.

Reproduce with:

python tools/compare_rdkit_mlx_conformers.py \
  --limit 20 \
  --n-conformers 10 \
  --out outputs/cheese_projection/rdkit_vs_mlx_conformers_20_k10_batched_scoring_after_fixes.csv

Architecture

Conformer Generation (N x k parallel)

SMILES x N
    |
[RDKit CPU] Extract params ONCE per molecule
    |
[Pack] SharedConstraintBatch (conf_to_mol indirection)
    |
+-- Stage 1: DG minimize (4D, Metal TPM=32) --------+
|   One threadgroup per conformer                    |
|   L-BFGS in-kernel, GPU-parallel line search       |
|   Shared constraints via conf_to_mol               |
+----------------------------------------------------+
    |
[Extract 3D] Drop 4th coordinate
    |
+-- Stage 2: ETK minimize (3D, Metal TPM=32) --------+
|   CSD torsion + improper + 1-2/1-3/1-4 distance     |
|   Optional parallel_grad for large molecules        |
+-----------------------------------------------------+
    |
+-- Stage 3: MMFF94 optimize (Metal, in-kernel) ------+
|   7 energy terms: bond, angle, stretch-bend,         |
|   OOP, torsion, vdW, electrostatic                   |
|   BFGS (default) or L-BFGS option                    |
+------------------------------------------------------+
    |
Optimized 3D conformers

Clustering (divide-and-conquer for 150k+)

Morgan FP (RDKit CPU)
        |
   uint8 -> uint32 packing
        |
+-- N <= 100k: Fused Metal Kernel ------+
|   Single dispatch, no NxN matrix       |
+-- N > 100k: Blockwise D&C ------------+
|   Tile both dimensions (auto-sized)    |
|   mx.eval() between tiles (free GPU)  |
+----------------------------------------+
        |
   Butina greedy (CPU, numpy CSR)
        |
   Clusters

Adaptive Iteration Scaling

Iterations auto-scale by molecule complexity (default). Small molecules converge early via in-kernel TOLX/gradient checks — no wasted GPU compute.

Formula: max_iters = base + scale * max(n_atoms, sqrt(n_constraints))

Molecule	Atoms	Constraints	DG iters	ETK iters	MMFF iters
Methane	5	10	400	200	275
Benzene	12	66	540	270	380
Aspirin	21	210	720	360	515
Testosterone	49	1176	1280	640	935
64-atom	64	2016	1580	790	1160

Override with explicit values when needed:

# Auto (default) — scales with molecule size
result = generate_conformers_nk(smiles_list, n_confs_per_mol=10)

# Fixed iterations for fine control
result = generate_conformers_nk(smiles_list, n_confs_per_mol=10,
    dg_max_iters=1000, etk_max_iters=500, mmff_max_iters=400)

Optimization Options

Option	Flag	Effect	When to use
Auto iterations (default)	`dg_max_iters=0`	Scales with molecule size	Always (default)
Warm-start retry	automatic	Re-runs non-converged with 2x iters	Always (automatic)
ETK parallel gradient	`parallel_grad=True`	1.18x ETK speedup	Many distance constraints
DG parallel gradient	`parallel_grad=True`	Parallelizes dist gradient	>500 distance constraints
MMFF94s variant	`mmff_variant="MMFF94s"`	Softer torsion barriers	Conjugated/aromatic molecules
MMFF L-BFGS	`mmff_use_lbfgs=True`	5x less memory	Molecules >50 atoms
MMFF BFGS (default)	`mmff_use_lbfgs=False`	2x faster for small mols	Molecules <50 atoms
ETKDG variant	`variant="ETKDGv3"`	8 variants supported	Choose per use case

MMFF94 Force Field Variants

Variant	Flag	Torsion barriers	Best for
MMFF94 (default)	`mmff_variant="MMFF94"`	Standard	General molecules
MMFF94s	`mmff_variant="MMFF94s"`	Softer for conjugated systems	Aromatic, planar, conjugated

Usage

3D Conformer Generation

from mlxmolkit.conformer_pipeline_v2 import generate_conformers_nk

# Basic: 10 conformers per molecule, ETKDGv2
result = generate_conformers_nk(
    smiles_list=["c1ccccc1", "CC(=O)O", "CC(=O)Oc1ccccc1C(=O)O"],
    n_confs_per_mol=10,
)
for mol in result.molecules:
    print(f"{mol.n_atoms} atoms, {len(mol.positions_3d)} conformers")

# Full pipeline with MMFF94
result = generate_conformers_nk(
    smiles_list=["c1ccccc1", "CC(=O)O"],
    n_confs_per_mol=50,
    variant="ETKDGv3",
    run_mmff=True,
    mmff_use_lbfgs=False,       # BFGS (default, fast for <50 atoms)
    max_confs_per_batch=500,    # divide-and-conquer batch size
)

# Pure distance geometry (no torsion refinement)
result = generate_conformers_nk(
    smiles_list=["c1ccccc1"],
    n_confs_per_mol=100,
    variant="DG",
)

# Large molecules: use L-BFGS for MMFF
result = generate_conformers_nk(
    smiles_list=[large_smiles],
    n_confs_per_mol=20,
    run_mmff=True,
    mmff_use_lbfgs=True,        # L-BFGS for >50 atoms
)

Example Script

# Basic: 20 molecules x 10 conformers
python examples/conf3d_example.py

# Scale test: 1000 molecules x 10 conformers with MMFF
python examples/conf3d_example.py --n-mols 1000 --n-confs 10 --mmff

# All options
python examples/conf3d_example.py --n-mols 100 --n-confs 20 --variant ETKDGv3 \
    --mmff --mmff-variant MMFF94s --batch-size 200

# Custom SMILES
python examples/conf3d_example.py --smiles "c1ccccc1" "CC(=O)O" --n-confs 50 --mmff

Molecular Clustering

from mlxmolkit import butina_tanimoto_mlx
import mlx.core as mx

# Automatic: fused kernel for N<=100k, blockwise for N>100k
result = butina_tanimoto_mlx(mx.array(fp_bytes), cutoff=0.4)
print(f"{len(result.clusters)} clusters")

Low-level API

from mlxmolkit import (
    fp_uint8_to_uint32,
    fused_neighbor_list_metal,
    tanimoto_neighbors_blockwise,
    butina_from_neighbor_list_csr,
)

fp_u32 = fp_uint8_to_uint32(mx.array(fp_bytes))

# Small N: fused single-dispatch
offsets, indices = fused_neighbor_list_metal(fp_u32, cutoff=0.4)

# Large N (150k+): divide-and-conquer blockwise
offsets, indices = tanimoto_neighbors_blockwise(fp_u32, cutoff=0.4)

result = butina_from_neighbor_list_csr(offsets, indices, N, cutoff=0.4)

Tests

pip install -e .
pytest tests/ -v

Conformer Quality vs RDKit

mlxmolkit conformers rescored by RDKit's MMFF94 for fair comparison (k=20, ETKDGv2):

After MMFF optimization, mlxmolkit conformers converge to the same energy basins as RDKit:

Molecule	Atoms	RMSD (pre-MMFF)	RMSD (post-MMFF)	E gap
Benzene	12	0.12 A	0.00 A	0.0
Aspirin	21	0.98 A	0.00 A	0.0
Ibuprofen	33	1.79 A	0.96 A	0.6
Acetaminophen	20	0.99 A	0.00 A	0.0

Full nvMolKit pipeline: DG (4D) → 4D→3D collapse → setReferenceValues → stereo checks → ETK (3D) → MMFF94. Bond/angle geometry matches RDKit within 0.03 A. After MMFF, energy gap is <1 kcal/mol.

References

nvMolKit — NVIDIA's CUDA implementation (Apache 2.0)
shivampatel10/mlxmolkit — TPM threadgroup kernels and MMFF Metal implementation
PYSEQM — LANL semi-empirical reference (BSD-3 Clause); the NumPy port in mlxmolkit/nddo/_pyseqm_port/ is a mechanical torch→numpy translation of selected modules
RDKit blog: Butina clustering with nvMolKit
MLX — Apple's ML framework with Metal kernel support
MMFF94 — Halgren, J. Comput. Chem. 1996
Butina, D. (1999) — Performance of Kier-Hall and molecular connectivity indices

Acknowledgements

Portions of the PM6_D / SCF / TETCI development were assisted by Claude (Anthropic). All commits are authored by the maintainer; Claude was used as a research/refactoring aide.

Name		Name	Last commit message	Last commit date
benchmarks		benchmarks
data		data
docs		docs
examples		examples
experiments		experiments
metal		metal
mlxmolkit		mlxmolkit
opencheese		opencheese
scripts		scripts
tests		tests
tools		tools
.gitignore		.gitignore
README.md		README.md
bench_bfgs.py		bench_bfgs.py
bench_compiled_vs_jit.py		bench_compiled_vs_jit.py
bench_conformers.py		bench_conformers.py
bench_head_to_head.py		bench_head_to_head.py
bench_mmff_100mol.py		bench_mmff_100mol.py
bench_multi_mol.py		bench_multi_mol.py
bench_native_vs_mlx.py		bench_native_vs_mlx.py
bench_profile_multi.py		bench_profile_multi.py
compare_rdkit_mlx.py		compare_rdkit_mlx.py
gfn2_gradient_aes_residual_report.md		gfn2_gradient_aes_residual_report.md
gxtb_retro_pseudocpp.md		gxtb_retro_pseudocpp.md
pyproject.toml		pyproject.toml
run_opt_comparison.py		run_opt_comparison.py
run_training_comparison.py		run_training_comparison.py
setup.py		setup.py

Navigation Menu

Repository navigation

About

Repository files navigation

mlxmolkit — GPU-accelerated molecular toolkit on Apple Silicon

What's new

Public API

Installation

Quick Start

Features

Performance

Conformer Generation (N=20 molecules, k=50 conformers = 1000 total)

Conformer Memory Scaling (DG + ETK + MMFF, batch=500)

Scale Tests

Batch Size Impact (N=20, k=50, C=1000)

GPU Memory per Conformer

Clustering (Enamine REAL subset, Apple M3 Max)

ETKDG Variant Comparison (N=20, k=50)

RDKit vs mlxmolkit conformer quality (N=20, k=10)

Architecture

Conformer Generation (N x k parallel)

Clustering (divide-and-conquer for 150k+)

Adaptive Iteration Scaling

Optimization Options

MMFF94 Force Field Variants

Usage

3D Conformer Generation

Example Script

Molecular Clustering

Low-level API

Tests

Conformer Quality vs RDKit

References

Acknowledgements

About

Releases

Packages

Contributors

Languages

Latest commit

History

Folders and files

Repository files navigation

mlxmolkit — GPU-accelerated molecular toolkit on Apple Silicon

What's new

Public API

Installation

Quick Start

Features

Performance

Conformer Generation (N=20 molecules, k=50 conformers = 1000 total)

Conformer Memory Scaling (DG + ETK + MMFF, batch=500)

Scale Tests

Batch Size Impact (N=20, k=50, C=1000)

GPU Memory per Conformer

Clustering (Enamine REAL subset, Apple M3 Max)

ETKDG Variant Comparison (N=20, k=50)

RDKit vs mlxmolkit conformer quality (N=20, k=10)

Architecture

Conformer Generation (N x k parallel)

Clustering (divide-and-conquer for 150k+)

Adaptive Iteration Scaling

Optimization Options

MMFF94 Force Field Variants

Usage

3D Conformer Generation

Example Script

Molecular Clustering

Low-level API

Tests

Conformer Quality vs RDKit

References

Acknowledgements

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages