# xvzemin/tace architecture

> Architecture diagram of the GitHub repository xvzemin/tace, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/xvzemin/tace
- Repository: https://github.com/xvzemin/tace

## Overview

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/models/_e3nn/tace.py`](https://github.com/xvzemin/tace/blob/main/tace/models/_e3nn/tace.py) (TACE models), [`tace/models/blocks.py`](https://github.com/xvzemin/tace/blob/main/tace/models/blocks.py) (Equivariant blocks), [`tace/dataset/datamodule.py`](https://github.com/xvzemin/tace/blob/main/tace/dataset/datamodule.py) (Data loading), [`tace/scripts/train.py`](https://github.com/xvzemin/tace/blob/main/tace/scripts/train.py) (Train command)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Data preparation"]
  node_read["Structure reading<br/>[read.py]"]
  node_graphdata["Graph construction<br/>[graph.py]"]
  node_neighbors["Neighbor lists<br/>[neighbour_list.py]"]
  node_datamodule["Data loading<br/>[datamodule.py]"]
  node_stats["Dataset statistics<br/>[statistics.py]"]
end

subgraph group_models["Potential models"]
  node_model["TACE models<br/>[tace.py]"]
  node_blocks["Equivariant blocks<br/>[blocks.py]"]
  node_radial["Radial features<br/>[radial.py]"]
  node_tece["TECE attention<br/>[interaction.py]"]
  node_scf["SCF models<br/>[model.py]"]
end

subgraph group_training["Training and evaluation"]
  node_traincmd["Train command<br/>[train.py]"]
  node_lightning["Training module<br/>[lit_model.py]"]
  node_trainer["Trainer<br/>[trainer.py]"]
  node_metrics["Property metrics<br/>[metrics_bake.py]"]
  node_finetune["Fine-tuning<br/>[finetune.py]"]
end

subgraph group_interfaces["Deployment interfaces"]
  node_ase["ASE calculator<br/>[calculator.py]"]
  node_lammps["LAMMPS interface<br/>[mliap.py]"]
  node_torchsim["TorchSim interface<br/>[torchsim.py]"]
  node_nvalchemi["NValCHEMI interface<br/>[wrapper.py]"]
  node_openmm["OpenMM integration<br/>[__init__.py]"]
  node_uspex["USPEX relaxation<br/>[relax.py]"]
end

subgraph group_engine["Equivariant engine"]
  node_eqxlinear["Equivariant layers<br/>[linear.py]"]
  node_eqxtensor["Tensor products<br/>[tensor_product.py]"]
  node_convolution["Equivariant convolutions<br/>[convolution.py]"]
  node_cuda["CUDA kernels<br/>[cuda.py]"]
end

node_researcher(("Model user"))
node_simulation(("Atomistic simulation"))

node_researcher -->|"starts training"| node_traincmd
node_traincmd -->|"loads data"| node_datamodule
node_read -->|"provides structures"| node_graphdata
node_graphdata -->|"forms neighborhoods"| node_neighbors
node_datamodule -->|"batches graphs"| node_graphdata
node_datamodule -->|"supports statistics"| node_stats
node_traincmd -->|"launches training"| node_trainer
node_trainer -->|"trains module"| node_lightning
node_lightning -->|"uses model"| node_model
node_lightning -->|"computes metrics"| node_metrics
node_finetune -->|"loads model"| node_model
node_model -->|"builds features"| node_blocks
node_model -->|"uses radial basis"| node_radial
node_model -.->|"SCF variant"| node_scf
node_model -.->|"TECE variant"| node_tece
node_blocks -.->|"equivariant maps"| node_eqxlinear
node_blocks -.->|"tensor couplings"| node_eqxtensor
node_blocks -.->|"convolution option"| node_convolution
node_convolution -.->|"accelerates kernels"| node_cuda
node_simulation -->|"requests predictions"| node_ase
node_ase -->|"evaluates structures"| node_model
node_simulation -->|"requests predictions"| node_lammps
node_simulation -->|"requests predictions"| node_torchsim
node_simulation -->|"requests predictions"| node_nvalchemi
node_openmm -.->|"routes through ASE"| node_ase
node_uspex -.->|"relaxes via LAMMPS"| node_lammps
node_lammps -->|"evaluates structures"| node_model
node_torchsim -->|"evaluates structures"| node_model
node_nvalchemi -->|"evaluates structures"| node_model

click node_traincmd "https://github.com/xvzemin/tace/blob/main/tace/scripts/train.py"
click node_read "https://github.com/xvzemin/tace/blob/main/tace/dataset/read.py"
click node_graphdata "https://github.com/xvzemin/tace/blob/main/tace/dataset/graph.py"
click node_neighbors "https://github.com/xvzemin/tace/blob/main/tace/dataset/neighbour_list.py"
click node_datamodule "https://github.com/xvzemin/tace/blob/main/tace/dataset/datamodule.py"
click node_stats "https://github.com/xvzemin/tace/blob/main/tace/dataset/statistics.py"
click node_model "https://github.com/xvzemin/tace/blob/main/tace/models/_e3nn/tace.py"
click node_blocks "https://github.com/xvzemin/tace/blob/main/tace/models/blocks.py"
click node_radial "https://github.com/xvzemin/tace/blob/main/tace/models/radial.py"
click node_tece "https://github.com/xvzemin/tace/blob/main/eqx/models/tace/tece_oam_rra/interaction.py"
click node_scf "https://github.com/xvzemin/tace/blob/main/tace/models/scf/model.py"
click node_lightning "https://github.com/xvzemin/tace/blob/main/tace/lightning/lit_model.py"
click node_trainer "https://github.com/xvzemin/tace/blob/main/tace/lightning/trainer.py"
click node_metrics "https://github.com/xvzemin/tace/blob/main/tace/utils/metrics_bake.py"
click node_finetune "https://github.com/xvzemin/tace/blob/main/tace/scripts/finetune.py"
click node_ase "https://github.com/xvzemin/tace/blob/main/tace/interface/ase/calculator.py"
click node_lammps "https://github.com/xvzemin/tace/blob/main/tace/interface/lammps/mliap.py"
click node_torchsim "https://github.com/xvzemin/tace/blob/main/tace/interface/torchsim/torchsim.py"
click node_nvalchemi "https://github.com/xvzemin/tace/blob/main/tace/interface/nvalchemi/wrapper.py"
click node_openmm "https://github.com/xvzemin/tace/blob/main/tace/interface/openmm/__init__.py"
click node_uspex "https://github.com/xvzemin/tace/blob/main/tace/interface/uspex26/relax.py"
click node_eqxlinear "https://github.com/xvzemin/tace/blob/main/eqx/o2/linear.py"
click node_eqxtensor "https://github.com/xvzemin/tace/blob/main/eqx/co3/tensor_product.py"
click node_convolution "https://github.com/xvzemin/tace/blob/main/eqx/conv/o3/convolution.py"
click node_cuda "https://github.com/xvzemin/tace/blob/main/eqx/kernels/cuda.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_read,node_graphdata,node_neighbors,node_datamodule,node_stats,node_researcher toneBlue
class node_model,node_blocks,node_radial,node_tece,node_scf toneAmber
class node_traincmd,node_lightning,node_trainer,node_metrics,node_finetune toneMint
class node_ase,node_lammps,node_torchsim,node_nvalchemi,node_openmm,node_uspex toneRose
class node_eqxlinear,node_eqxtensor,node_convolution,node_cuda,node_simulation toneIndigo
```

## Components

### Data preparation

- [Structure reading](https://github.com/xvzemin/tace/blob/main/tace/dataset/read.py): `tace/dataset/read.py`
- [Graph construction](https://github.com/xvzemin/tace/blob/main/tace/dataset/graph.py): `tace/dataset/graph.py`
- [Neighbor lists](https://github.com/xvzemin/tace/blob/main/tace/dataset/neighbour_list.py): `tace/dataset/neighbour_list.py`
- [Data loading](https://github.com/xvzemin/tace/blob/main/tace/dataset/datamodule.py): `tace/dataset/datamodule.py`
- [Dataset statistics](https://github.com/xvzemin/tace/blob/main/tace/dataset/statistics.py): `tace/dataset/statistics.py`

### Potential models

- [TACE models](https://github.com/xvzemin/tace/blob/main/tace/models/_e3nn/tace.py): `tace/models/_e3nn/tace.py`
- [Equivariant blocks](https://github.com/xvzemin/tace/blob/main/tace/models/blocks.py): `tace/models/blocks.py`
- [Radial features](https://github.com/xvzemin/tace/blob/main/tace/models/radial.py): `tace/models/radial.py`
- [TECE attention](https://github.com/xvzemin/tace/blob/main/eqx/models/tace/tece_oam_rra/interaction.py): `eqx/models/tace/tece_oam_rra/interaction.py`
- [SCF models](https://github.com/xvzemin/tace/blob/main/tace/models/scf/model.py): `tace/models/scf/model.py`

### Training and evaluation

- [Train command](https://github.com/xvzemin/tace/blob/main/tace/scripts/train.py): `tace/scripts/train.py`
- [Training module](https://github.com/xvzemin/tace/blob/main/tace/lightning/lit_model.py): `tace/lightning/lit_model.py`
- [Trainer](https://github.com/xvzemin/tace/blob/main/tace/lightning/trainer.py): `tace/lightning/trainer.py`
- [Property metrics](https://github.com/xvzemin/tace/blob/main/tace/utils/metrics_bake.py): `tace/utils/metrics_bake.py`
- [Fine-tuning](https://github.com/xvzemin/tace/blob/main/tace/scripts/finetune.py): `tace/scripts/finetune.py`

### Deployment interfaces

- [ASE calculator](https://github.com/xvzemin/tace/blob/main/tace/interface/ase/calculator.py): `tace/interface/ase/calculator.py`
- [LAMMPS interface](https://github.com/xvzemin/tace/blob/main/tace/interface/lammps/mliap.py): `tace/interface/lammps/mliap.py`
- [TorchSim interface](https://github.com/xvzemin/tace/blob/main/tace/interface/torchsim/torchsim.py): `tace/interface/torchsim/torchsim.py`
- [NValCHEMI interface](https://github.com/xvzemin/tace/blob/main/tace/interface/nvalchemi/wrapper.py): `tace/interface/nvalchemi/wrapper.py`
- [OpenMM integration](https://github.com/xvzemin/tace/blob/main/tace/interface/openmm/__init__.py): `tace/interface/openmm/__init__.py`
- [USPEX relaxation](https://github.com/xvzemin/tace/blob/main/tace/interface/uspex26/relax.py): `tace/interface/uspex26/relax.py`

### Equivariant engine

- [Equivariant layers](https://github.com/xvzemin/tace/blob/main/eqx/o2/linear.py): `eqx/o2/linear.py`
- [Tensor products](https://github.com/xvzemin/tace/blob/main/eqx/co3/tensor_product.py): `eqx/co3/tensor_product.py`
- [Equivariant convolutions](https://github.com/xvzemin/tace/blob/main/eqx/conv/o3/convolution.py): `eqx/conv/o3/convolution.py`
- [CUDA kernels](https://github.com/xvzemin/tace/blob/main/eqx/kernels/cuda.py): `eqx/kernels/cuda.py`

### Other components

- **Model user**
- **Atomistic simulation**

## Connections

- 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