CEITNet: Cartesian Environment Interaction Tensor Network

CEITNet is a network for high-order tensor prediction.

Installation

# Create the Conda env:
conda create -n ceitnet python=3.13 -y
conda activate ceitnet
# Install PyTorch 2.8 and torch_scatter according to your local CPU/CUDA setup:
pip install torch==2.8.0 
pip install torch-scatter -f https://data.pyg.org/whl/torch-2.6.0+${CUDA}.html
# Install the remaining deps:
pip install torch_geometric pymatgen jarvis-tools e3nn pandas spglib=2.6.0 tqdm pandarallel

Datasets

The files in repro_data/preprocessed_data/ are compressed. You can use it via xz -dk [XZ_DATA].

We recommend constructing the preprocessed data from the original preprocessing pipeline: See GMTNet repository for more information.

Usage

To test using the pretrained checkpoints:

python [TASK]/test.py --zero_mask

To train the model from scratch for a specific task:

python [TASK]/train.py

To test a trained model using a specific checkpoint:

python [TASK]/test.py --ckpt_path [CKPT_PATH] --zero_mask

• --ckpt_path: Path to the model checkpoint.
• --zero_mask: (Optional) Enforces zero-component constraints. Derived from GMTNet.

Notes

This codebase is primarily designed for the included benchmarks. The training pipeline follows GMTNet, including the hyperparameters, optimization settings, and data-splitting protocols. We cleaned up the original codebase. Performance may vary with different configurations, and additional engineering may be needed for larger-scale systems or new tensorial prediction tasks.

Acknowledgement

The training pipeline are derived from GMTNet. We thank the authors for their great work and open-source contribution.