MolNex is a molecular machine-learning framework organized as four cooperating packages — molix, molrep, molpot, and molzoo — that can be used together or independently.

Under active development. Public APIs may change between minor releases.

Vision

Molecular ML projects rarely need just one model. They accumulate training infrastructure, representation modules, compositional output layers, and several reference architectures — and when all of that collapses into a single library built around one preferred model, the codebase becomes hard to extend, test, and reason about.

MolNex exists to keep those concerns apart so they can grow at their own pace. It aspires to be a framework where a contributor always knows which layer they are working in, what that layer is allowed to own, and how it connects to the rest of the stack — rather than a trainer wrapped around one favorite architecture.

That separation is what unlocks the goal: multiple model families living in one codebase, training infrastructure and representation modules reused freely across them, and reference models that stay replaceable examples of assembly instead of becoming the framework itself.

Capabilities

Package	Capability
molix	Training and execution — Trainer, TrainState, step protocol, hook lifecycle, packed-cache data pipeline, dataset loaders (QM9, RevMD17, 3BPA, Water-LES, MolRec), checkpointing, torch.compile / CUDA-graph and AOT-Inductor export, an in-process Langevin MD driver, trajectory diagnostics, and native C++ ops
molrep	Representation learning — node/radial/angular embeddings, cutoffs, equivariant interaction blocks (tensor products, symmetric contraction), pooling and readout heads
molpot	Potentials and composition — classical potential terms (LJ, harmonic bonds/angles/dihedrals, electrostatics), prediction heads, force/stress derivation, and the PotentialComposer assembly layer
molzoo	Assembled reference models — curated encoder and potential families (MACE, Allegro, PiNet, Sonata), each with a paper-traceable spec

Install

pip install molnex

Requires Python >= 3.10 and PyTorch >= 2.10. The package builds native C++ ops via scikit-build-core and CMake >= 4.0; an editable install is pip install -e ".[dev]". See Installation for the full build setup.

Quick start

Train a model with molix.Trainer:

import torch
from molix.core.trainer import Trainer

trainer = Trainer(
    model=model,                       # any nn.Module
    loss_fn=loss_fn,                   # (predictions, batch) -> scalar
    optimizer_factory=lambda p: torch.optim.SGD(p, lr=0.1),
)

state = trainer.train(datamodule, max_epochs=5)
print(state["train/loss"])

See the Molix Quick Start for the runnable end-to-end version, and Train a Graph Model for molecular graph batches.

Documentation

• Documentation home
• Installation
• Molix — training, hooks, data, and execution
• MolRep — representation learning modules
• MolPot — potential composition and physical outputs
• MolZoo — reference encoder families
• API Reference

MolCrafts ecosystem

Project	Role
molpy	Python toolkit — the shared molecular data model & workflow layer
molrs	Rust core — molecular data structures & compute kernels (native + WASM)
molpack	Packmol-grade molecular packing (Rust + Python)
molvis	WebGL molecular visualization & editing
molexp	Workflow & experiment-management platform
molnex	Molecular machine-learning framework — this repo
molq	Unified job queue — local / SLURM / PBS / LSF
molcfg	Layered configuration library
mollog	Structured logging, stdlib-compatible
molhub	Molecular dataset hub
molmcp	MCP server for the ecosystem
molrec	Atomistic record specification

Contributing

Contributions are welcome — see the documentation to get started.

License

BSD-3-Clause — see LICENSE.

---