SpinGTP

Official codebase for Spin-Weighted Spherical Harmonics Enable Complete and Scalable E(3)-Equivariant Networks

Spherical Harmonics

Spin-Weighted Spherical Harmonics

Overview

SpinGTP provides the core implementation of spin-weighted spherical harmonic representations and tensor-product operations for scalable E(3)-equivariant neural networks.

Codebase Structure

ROOT
├── 3BPA/                            Experiments on the 3BPA dataset
├── fairchem/                        Experiments on OC20 using the FairChem codebase
├── spice-mace-off/                  Experiments on SPICE-MACE-OFF chiral subset
├── spingtp/                         Core SpinGTP codebase, including SWSH irreps and tensor-product algorithms
├── tetris/                          Experiments on the Tetris dataset
├── LICENCE                          License information
├── README.md                        Project overview and usage instructions
├── sphharm_professional_whole.gif   Spherical harmonics visualization
└── swsh_professional_14cols.gif     SWSH visualization

Quick Install

Install the core SpinGTP package

python -m pip install -e spingtp

Implementaion of SpinGTP

SpinGTP is organized around a core SWSH package plus model integrations.

The representation layer lives in spingtp/spingtp/modules/swsh_irreps.py, swsh.py, and swsh_reference.py. It provides SWSHIrrep, SWSHIrreps, SWSHSphericalHarmonics, frame builders such as frame_from_node_quadrupole and frame_from_global_position.

The equivariant network layers live in swsh_tensor_product.py, swsh_linear.py, swsh_activation.py, swsh_gate.py, swsh_layer_norm.py, swsh_symmetric_contraction.py, and swsh_spin_head.py. These files implement the main tensor-product backends (SWSHTensorProduct, FullyConnectedSWSHTensorProduct, DepthwiseSWSHTensorProduct), channel mixing (SWSHLinear), scalar activations and gates, equivariant normalization, and symmetric contractions.

Quick Examples

import torch
from spingtp.modules.swsh import SWSHSphericalHarmonics
from spingtp.modules.swsh_irreps import SWSHIrreps

irreps = SWSHIrreps("1x(0,0)e+1x(1,0)e+1x(1,1)o")
edge_vec = torch.randn(8, 3)

swsh = SWSHSphericalHarmonics(irreps)
edge_attr = swsh(edge_vec)
print(edge_attr.shape)

Acknowledgement

This work was supported in part by the National Science Foundation under Grants IIS-2243850, CNS-2328395, and MOMS-2331036; the National Institutes of Health under Grant U01AG070112; the Texas A&M University Division of Research Targeted Proposal Teams Funding Program; and the Texas A&M Institute of Data Science Thematic Labs Program.

Citation