How xvzemin/tace works
TACE is a machine-learning interatomic-potential toolkit built around tensor atomic and edge cluster expansion models. Its main workflow is to read and prepare atomistic structures, train or fine-tune an equivariant model against requested properties, then evaluate or deploy predictions through simulation interfaces. The graph also preserves documented model variants, physical-property/plugin capabilities, and EquivariantX’s lower-level equivariant and accelerated-computation layer. Source excerpts are partial, so uncertain internal call wiring is omitted; documented capabilities are shown without asserting uninspected implementation details.
Main entry points: tace.py TACE models · blocks.py Equivariant blocks · datamodule.py Data loading · train.py Train command
Components 27
Data preparation
- Structure reading
tace/dataset/read.py - Graph construction
tace/dataset/graph.py - Neighbor lists
tace/dataset/neighbour_list.py - Data loading
tace/dataset/datamodule.py - Dataset statistics
tace/dataset/statistics.py
Potential models
- TACE models
tace/models/_e3nn/tace.py - Equivariant blocks
tace/models/blocks.py - Radial features
tace/models/radial.py - TECE attention
eqx/models/tace/tece_oam_rra/interaction.py - SCF models
tace/models/scf/model.py
Training and evaluation
- Train command
tace/scripts/train.py - Training module
tace/lightning/lit_model.py - Trainer
tace/lightning/trainer.py - Property metrics
tace/utils/metrics_bake.py - Fine-tuning
tace/scripts/finetune.py
Deployment interfaces
- ASE calculator
tace/interface/ase/calculator.py - LAMMPS interface
tace/interface/lammps/mliap.py - TorchSim interface
tace/interface/torchsim/torchsim.py - NValCHEMI interface
tace/interface/nvalchemi/wrapper.py - OpenMM integration
tace/interface/openmm/__init__.py - USPEX relaxation
tace/interface/uspex26/relax.py
Equivariant engine
- Equivariant layers
eqx/o2/linear.py - Tensor products
eqx/co3/tensor_product.py - Equivariant convolutions
eqx/conv/o3/convolution.py - CUDA kernels
eqx/kernels/cuda.py
Other components
- Model user
- Atomistic simulation
Connections 29
- Model user → Train command: starts training
- Train command → Data loading: loads data
- Structure reading → Graph construction: provides structures
- Graph construction → Neighbor lists: forms neighborhoods
- Data loading → Graph construction: batches graphs
- Data loading → Dataset statistics: supports statistics
- Train command → Trainer: launches training
- Trainer → Training module: trains module
- Training module → TACE models: uses model
- Training module → Property metrics: computes metrics
- Fine-tuning → TACE models: loads model
- TACE models → Equivariant blocks: builds features
- TACE models → Radial features: uses radial basis
- TACE models → SCF models: SCF variant
- TACE models → TECE attention: TECE variant
- Equivariant blocks → Equivariant layers: equivariant maps
- Equivariant blocks → Tensor products: tensor couplings
- Equivariant blocks → Equivariant convolutions: convolution option
- Equivariant convolutions → CUDA kernels: accelerates kernels
- Atomistic simulation → ASE calculator: requests predictions
- ASE calculator → TACE models: evaluates structures
- Atomistic simulation → LAMMPS interface: requests predictions
- Atomistic simulation → TorchSim interface: requests predictions
- Atomistic simulation → NValCHEMI interface: requests predictions
- OpenMM integration → ASE calculator: routes through ASE
- USPEX relaxation → LAMMPS interface: relaxes via LAMMPS
- LAMMPS interface → TACE models: evaluates structures
- TorchSim interface → TACE models: evaluates structures
- NValCHEMI interface → TACE models: evaluates structures