Introduction

SO(3) is the group of three dimensional rotations. SO(3)-equivariant neural networks learn functions that behave in specific ways with respect to this group. Such architectures have recently become popular in a range of applications in physics, chemistry, computational biology and computer vision.

From a mathematical point of view, the key feature of these architectures is that the output of each neuron is an SO(3)-covariant quantity, which means that under rotations it transforms either according to a single irreducible representation (irrep) of SO(3) or a specific combination of irreducible representations. This property is essential for ensuring equivariance, but also puts restrictions on what operations are permissible in the network.

GElib provides two basic classes to facilitate implementing SO(3)-equivariant networks:

  1. The SO3part class stores any quantity that transforms according to a single irrep.

  2. The SO3vec class stores a quantity that transforms according to a combination of different irreps.

The library also implements the most important operations that can be performed on SO(3)-covariant objects, specifically, different variants of the Clebsch–Gordan product.

Features

GElib provides the following features:

  1. Classes to store and manipulate SO(3)-covariant (equivariant) vectors.

  2. Fast implementations of different variants of the Clebsch-Gordan product on both the CPU and the GPU.

  3. Facilities for operating on arrays of SO(3)-covariant vectors in parallel, including in irregular patterns (graphs).

  4. Interoperability with PyTorch’s autograd functionality.

Installation

GElib is installed as a PyTorch C++ (or CUDA) extension and requires the following:

  1. C++17 or higher

  2. PyTorch

  3. cnine (see below)

GElib is easiest to install with pip:

  1. Download the cnine tensor library and install it on your system

  2. Download GElib. By default, it is assumed that cnine and GElib are downloaded to the same directory (e.g., Downloads).

  3. Edit the user configurable variables in python/setup.py as necessary.

  4. Run the command pip install -e . in the GElib/python directory.

Configuration

The installation can be configured by setting the following variables in python/setup.py.

compile_with_cuda

If set to True, GElib will be compiled with GPU suport. This requires a working CUDA and CUBLAS installation on your system and PyTorch itself having been compiled with CUDA enabled. If GElib is compiled with CUDA, you must always import torch before importing GElib.

copy_warnings

If set to True, GElib will print a message to the terminal whenever a data object is copied or move-copied. This option is useful for code optimization.

torch_convert_warnings

If set to True, cnine will print a message to the terminal whenver a data object is explicitly converted (as opposed to just forming a tensor view) to/from PyTorch format.

Troubleshooting

  1. PyTorch C++ extensions must be compiled against the same version of CUDA that PyTorch itself was compiled with. If this becomes an issue, it might be necessary to install an alternative version of CUDA on your system and force setuptools to use that version by setting the CUDA_HOME enironment variable, as, e.g.

    export CUDA_HOME=/usr/local/cuda-11.3

Known issues

GElib is undergoing a major version change. Not all of its functionality is currently live.

Usage

To load the gelib module in Python, use the command import gelib. Depending on whether or not GElib was installed with CUDA, the message Starting GElib with GPU support or Starting GElib without GPU support will appear, confirming that gelib has successfully started and initialized its static data objects.