add pos-egnn

models/fm4m.py

@@ -32,6 +32,7 @@ from models.selfies_ted.load import SELFIES as bart
 from models.mhg_model import load as mhg
 from models.smi_ted.smi_ted_light.load import load_smi_ted
 from models.smi_ssed.load import load_smi_ssed
+from models.pos_egnn.load import POSEGNN as pos
 
 import mordred
 from mordred import Calculator, descriptors
@@ -60,6 +61,7 @@ def avail_models_data():
   {"Name": "mol-xl","Model Name": "MolFormer", "Description": "MolFormer model for string based SMILES modality", "Timestamp": "2024-06-21 12:35:56"},
   {"Name": "mhg", "Model Name": "MHG-GED","Description": "Molecular hypergraph model", "Timestamp": "2024-07-10 00:09:42"},
   {"Name": "smi-ted", "Model Name": "SMI-TED","Description": "SMILES based encoder decoder model", "Timestamp": "2024-07-10 00:09:42"},
+  {"Name": "pos-egnn", "Model Name": "POS-EGNN","Description": "3D atom position model", "Timestamp": "2025-04-04 00:11:42"},
   {"Name": "smi-ssed", "Model Name": "SMI-SSED","Description": "SMILES based encoder decoder model", "Timestamp": "2024-07-10 00:09:42"}]
 
 
@@ -71,6 +73,7 @@ def avail_models(raw=False):
               {"Name": "bart","Model Name": "SELFIES-TED","Description": "BART model for string based SELFIES modality"},
               {"Name": "mol-xl","Model Name": "MolFormer", "Description": "MolFormer model for string based SMILES modality"},
               {"Name": "mhg", "Model Name": "MHG-GED","Description": "Molecular hypergraph model"},
+              {"Name": "pos", "Model Name": "POS-EGNN","Description": "3D atom position model"},
               {"Name": "Mordred", "Model Name": "Mordred","Description": "Baseline: A descriptor-calculation software application that can calculate more than 1800 two- and three-dimensional descriptors"},
               {"Name": "MorganFingerprint", "Model Name": "MorganFingerprint","Description": "Baseline: Circular atom environments based descriptor"}              
   ]
@@ -150,6 +153,7 @@ def reset():
      {"Name": "mol-xl", "Description": "MolFormer model for string based SMILES modality",
       "Timestamp": "2024-06-21 12:35:56"},
      {"Name": "mhg", "Description": "MHG", "Timestamp": "2024-07-10 00:09:42"},
+     {"Name": "pos", "Description": "POS-EGNN", "Timestamp": "2024-07-10 00:09:42"},
      {"Name": "spec-gru", "Description": "Spectrum modality with GRU", "Timestamp": "2024-07-10 00:09:42"},
      {"Name": "spec-lstm", "Description": "Spectrum modality with LSTM", "Timestamp": "2024-07-10 00:09:54"},
      {"Name": "3d-vae", "Description": "VAE model for 3D atom positions", "Timestamp": "2024-07-10 00:10:08"}]
@@ -204,7 +208,9 @@ avail_models_data()
 
 
 def get_representation(train_data,test_data,model_type, return_tensor=True):
-    alias = {"MHG-GED": "mhg", "SELFIES-TED": "bart", "MolFormer": "mol-xl", "Molformer": "mol-xl", "SMI-TED": "smi-ted", "SMI-SSED": "smi-ssed"}
+    #alias = {"MHG-GED": "mhg", "SELFIES-TED": "bart", "MolFormer": "mol-xl", "Molformer": "mol-xl", "SMI-TED": "smi-ted", "SMI-SSED": "smi-ssed"}
+    alias = {"MHG-GED": "mhg", "SELFIES-TED": "bart", "SMI-TED": "smi-ted", "POS-EGNN": "pos", "MolFormer": "mol-xl", "Molformer": "mol-xl",
+             }
     if model_type in alias.keys():
         model_type = alias[model_type]
 
@@ -227,6 +233,12 @@ def get_representation(train_data,test_data,model_type, return_tensor=True):
         x_batch = model.encode(train_data, return_tensor=return_tensor)
         x_batch_test = model.encode(test_data, return_tensor=return_tensor)
 
+    elif model_type == "pos":
+        model = pos()
+        model.load()
+        x_batch = model.encode(train_data, return_tensor=return_tensor)
+        x_batch_test = model.encode(test_data, return_tensor=return_tensor)
+
     elif model_type == "smi-ted":
         model = load_smi_ted(folder='../models/smi_ted/smi_ted_light', ckpt_filename='smi-ted-Light_40.pt')
         with torch.no_grad():
@@ -314,7 +326,9 @@ def get_representation(train_data,test_data,model_type, return_tensor=True):
 
 def single_modal(model,dataset=None, downstream_model=None, params=None, x_train=None, x_test=None, y_train=None, y_test=None):
     print(model)
-    alias = {"MHG-GED":"mhg", "SELFIES-TED": "bart", "MolFormer":"mol-xl", "Molformer": "mol-xl", "SMI-TED": "smi-ted", "SMI-SSED": "smi-ssed"}
+    #alias = {"MHG-GED":"mhg", "SELFIES-TED": "bart", "MolFormer":"mol-xl", "Molformer": "mol-xl", "SMI-TED": "smi-ted", "SMI-SSED": "smi-ssed"}
+    alias = {"MHG-GED": "mhg", "SELFIES-TED": "bart",  "SMI-TED": "smi-ted", "POS-EGNN": "pos", "MolFormer": "mol-xl", "Molformer": "mol-xl",
+             }
     data = avail_models(raw=True)
     df = pd.DataFrame(data)
     #print(list(df["Name"].values))
@@ -383,13 +397,14 @@ def single_modal(model,dataset=None, downstream_model=None, params=None, x_train
         print(f'x_batch_test Nan index: {nan_indices}')
         print(f'x_batch_test shape: {x_batch_test.shape}, y_batch_test len: {len(y_batch_test)}')
 
-    print(f" Calculating ROC AUC Score ...")
+    print(f"Model selected: {downstream_model} - Calculating ROC AUC Score ...")
 
     if downstream_model == "XGBClassifier":
         if params == None:
             xgb_predict_concat = XGBClassifier()
         else:
             xgb_predict_concat = XGBClassifier(**params) # n_estimators=5000, learning_rate=0.01, max_depth=10
+
         xgb_predict_concat.fit(x_batch, y_batch)
 
         y_prob = xgb_predict_concat.predict_proba(x_batch_test)[:, 1]
@@ -628,7 +643,10 @@ def multi_modal(model_list,dataset=None, downstream_model=None,params=None, x_tr
     df = pd.DataFrame(data)
     list(df["Name"].values)
 
-    alias = {"MHG-GED":"mhg", "SELFIES-TED": "bart", "MolFormer":"mol-xl",  "Molformer": "mol-xl","SMI-TED":"smi-ted","SMI-SSED":"smi-ssed", "Mordred": "Mordred", "MorganFingerprint": "MorganFingerprint"}
+    #alias = {"MHG-GED":"mhg", "SELFIES-TED": "bart", "MolFormer":"mol-xl",  "Molformer": "mol-xl","SMI-TED":"smi-ted","SMI-SSED":"smi-ssed", "Mordred": "Mordred", "MorganFingerprint": "MorganFingerprint"}
+    alias = {"MHG-GED": "mhg", "SELFIES-TED": "bart", "SMI-TED": "smi-ted", "POS-EGNN": "pos", "MolFormer": "mol-xl",
+             "Molformer": "mol-xl", "Mordred": "Mordred", "MorganFingerprint": "MorganFingerprint"
+             }
     #if set(model_list).issubset(list(df["Name"].values)):
     if set(model_list).issubset(list(alias.keys())):
         for i, model in enumerate(model_list):
@@ -717,7 +735,7 @@ def multi_modal(model_list,dataset=None, downstream_model=None,params=None, x_tr
 
         print("Generating latent plots : Done")
 
-    print(f" Calculating ROC AUC Score ...")
+    print(f"Model selected: {downstream_model} - Calculating ROC AUC Score ...")
 
 
     if downstream_model == "XGBClassifier":

models/pos_egnn/.gitignore

@@ -0,0 +1,158 @@
+data/
+!/morningstar/data/
+!/tests/data
+.ruff*
+lightning_logs/
+.DS_Store
+.vscode
+
+# Useful scratch subfolder
+_scratch
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+*.Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+temp_data/
+.aim/
+csv_logs/
+
+# Dataset files
+*data.mdb
+*.zip
+*.tar.gz
+
+/plots
+/checkpoints
+/data*
+
+*.err
+*.out
+*.ckpt
+output.xyz

models/pos_egnn/README.md

@@ -0,0 +1,40 @@
+# Position-based Equivariant Graph Neural Network (`pos-egnn`)
+This repository contains PyTorch source code for loading and performing inference using the `pos-egnn`, a foundation model for Chemistry and Materials.
+
+**GitHub**: https://github.com/ibm/materials
+
+**HuggingFace**: https://huggingface.co/ibm-research/materials.pos-egnn
+
+<p align="center">
+    <img src="../../img/posegnn.svg">
+</p>
+
+## Introduction
+We present `pos-egnn`, a Position-based Equivariant Graph Neural Network foundation model for Chemistry and Materials. The model was pre-trained on 1.4M samples (i.e., 90%) from the Materials Project Trajectory (MPtrj) dataset to predict energies, forces and stress. `pos-egnn` can be used as a machine-learning potential, as a feature extractor, or can be fine-tuned for specific downstream tasks.
+
+Besides the model weigths `pos-egnn.v1-6M.pt` (download from [HuggingFace](https://huggingface.co/ibm-research/materials.pos-egnn)), we also provide an `example.ipynb` notebook (download from [GitHub](https://github.com/ibm/materials)), which demonstrates how to perform inference, feature extraction and molecular dynamics simulation with the model.
+
+For more information, please reach out to [email protected] and/or [email protected]
+
+## Table of Contents
+1. [**Getting Started**](#getting-started)
+2. [**Example**](#example)
+
+## Getting Started
+Follow these steps to replicate our environment and install the necessary libraries:
+
+First, make sure to have Python 3.11 installed. Then, to create the virtual environment, run the following commands:
+
+```bash
+python3.11 -m venv env
+source env/bin/activate
+```
+
+Run  the following command to install the library dependencies.
+
+```bash
+pip install -r requirements.txt
+```
+
+## Example
+Please refer to the `example.ipynb` for a step-by-step demonstration on how to perform inference, feature extraction and molecular dynamics simulation with the model.

models/pos_egnn/example.ipynb

@@ -0,0 +1,250 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# POS-EGNN "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Setup"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Uncomment to install notebook-only dependencies\n",
+    "# !pip install nglview ipywidgets"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4bac12c6048044898065f0778d95caeb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": []
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import nglview as nv\n",
+    "import torch\n",
+    "from ase import units\n",
+    "from ase.io import read\n",
+    "from ase.md.langevin import Langevin\n",
+    "\n",
+    "from posegnn.calculator import PosEGNNCalculator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = \"cpu\"\n",
+    "torch.set_float32_matmul_precision(\"high\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Feature Extraction"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Please download checkpoint from https://huggingface.co/ibm-research/materials.pos-egnn\n",
+    "calculator = PosEGNNCalculator(\"pos-egnn.v1-6M.ckpt\", device=device, compute_stress=False)\n",
+    "atoms = read(\"inputs/3BPA.xyz\", index=0)\n",
+    "atoms.calc = calculator"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "torch.Size([27, 256])"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "embeddings = atoms.get_invariant_embeddings()\n",
+    "embeddings.shape"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "energy = atoms.get_potential_energy()\n",
+    "forces = atoms.get_forces()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "(array([-175.05188], dtype=float32),\n",
+       " array([[ 0.34280202, -0.41967863,  0.7246248 ],\n",
+       "        [-0.86854756, -0.12186409, -2.305024  ],\n",
+       "        [ 0.26306945,  0.06607065,  0.85476065],\n",
+       "        [-0.230737  ,  0.02304646, -0.5161394 ],\n",
+       "        [-0.43901953,  2.7678285 , -0.70297724],\n",
+       "        [ 0.03933215, -0.50390136,  1.0451801 ],\n",
+       "        [ 0.37628424, -2.2708364 , -0.7662437 ],\n",
+       "        [ 0.25884533, -1.6086004 , -0.08700082],\n",
+       "        [-0.09319548, -0.24666801, -0.48069426],\n",
+       "        [ 0.01849201,  1.001767  ,  2.151208  ],\n",
+       "        [-0.46055827,  1.3630681 , -0.38470453],\n",
+       "        [ 0.38605827, -0.32170498,  0.6269282 ],\n",
+       "        [-0.29103595,  0.22509174, -0.26729944],\n",
+       "        [ 1.3340423 , -1.727819  , -0.08812339],\n",
+       "        [-0.96442086,  1.1447092 ,  1.0665402 ],\n",
+       "        [-0.74679977,  0.56782806,  0.03098067],\n",
+       "        [ 0.42040402,  0.7405614 , -0.6953748 ],\n",
+       "        [-0.25654212,  0.25282693,  0.25414664],\n",
+       "        [ 2.0051584 , -0.38257334, -0.26911467],\n",
+       "        [-0.00743119,  0.43786597, -0.27683535],\n",
+       "        [ 0.64563835, -0.5602143 , -0.11240276],\n",
+       "        [-0.00601408, -1.03808   ,  0.23635206],\n",
+       "        [-0.04149596,  0.02955294, -0.06748012],\n",
+       "        [-0.86066115,  0.00299245,  0.06783121],\n",
+       "        [-0.05461264,  0.05352221, -0.06844339],\n",
+       "        [-0.26291835,  0.58347785,  0.19614606],\n",
+       "        [-0.50613666, -0.05826864, -0.16684091]], dtype=float32))"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "energy, forces"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Molecular Dynamics Simulation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "True"
+      ]
+     },
+     "execution_count": 11,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "dyn = Langevin(atoms=atoms, friction=0.005, temperature_K=310, timestep=0.5 * units.fs)\n",
+    "\n",
+    "def write_frame():\n",
+    "    dyn.atoms.write(\"output.xyz\", append=True)\n",
+    "\n",
+    "dyn.attach(write_frame, interval=5)\n",
+    "dyn.run(500)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "139d7605baca43d79ea515d3454d9941",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "NGLWidget(max_frame=234)"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "traj = read('output.xyz', index=slice(None))\n",
+    "view = nv.show_asetraj(traj)\n",
+    "display(view)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "py311",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

models/pos_egnn/inputs/3BPA.xyz

@@ -0,0 +1,29 @@
+27
+Lattice="50.0 0.0 0.0 0.0 50.0 0.0 0.0 0.0 50.0" Properties=species:S:1:pos:R:3
+C        0.32656990      -1.01286015       0.72107275
+C       -0.19461567       0.25830309       0.83413890
+C        0.79249430      -1.60244652       1.90438013
+H        0.43514812      -1.59784402      -0.17897327
+C       -0.06984110       0.97877174       2.04804512
+O       -0.82551374       0.94764149      -0.23208227
+N       -0.29006580       2.37382353       1.86234713
+N        0.38837699       0.47664821       3.17063941
+H       -0.98416331       2.63164516       1.17783211
+H       -0.35613357       2.91935548       2.65951281
+C        0.72925588      -0.88036074       3.09652633
+H        0.89111308      -1.33212550       4.04892841
+H        1.21249480      -2.61699312       1.95761271
+C       -1.55059042       0.13992010      -1.19894310
+C       -2.93143480      -0.31868852      -0.84882155
+H       -0.88027102      -0.72270141      -1.50495635
+H       -1.71075243       0.76854244      -2.05655313
+C       -3.36709795      -1.63868610      -0.72393961
+C       -3.88193920       0.73797993      -0.59066374
+C       -4.66143774      -1.92714685      -0.27081530
+H       -2.71536973      -2.46157596      -0.92494411
+C       -5.57704224      -0.87827710      -0.11327951
+H       -4.95294558      -2.96156288      -0.08227495
+C       -5.17381471       0.43421272      -0.29908870
+H       -6.59525001      -1.10219470       0.20624618
+H       -5.92759350       1.19353463      -0.20195647
+H       -3.44855719       1.74786696      -0.56747187

models/pos_egnn/load.py

@@ -0,0 +1,63 @@
+import torch
+from .posegnn.calculator import PosEGNNCalculator
+import ase
+from ase import Atoms
+from rdkit import Chem
+from rdkit.Chem import AllChem
+import pandas as pd
+import numpy as np
+from huggingface_hub import hf_hub_download
+from tqdm import tqdm
+
+torch.set_float32_matmul_precision("high")
+
+def smiles_to_atoms(smiles):
+    mol = Chem.AddHs(Chem.MolFromSmiles(smiles))
+    AllChem.EmbedMolecule(mol)
+    ase_atoms = ase.Atoms(
+        numbers=[
+            atom.GetAtomicNum() for atom in mol.GetAtoms()
+        ],
+        positions=mol.GetConformer().GetPositions()
+    )
+    return ase_atoms
+
+class POSEGNN():
+    def __init__(self, use_gpu=True):
+        device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
+        self.device = device
+        self.calculator = None
+
+    def load(self, checkpoint=None):
+        repo_id = "ibm-research/materials.pos-egnn"
+        filename = "pytorch_model.bin"
+        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+        self.calculator = PosEGNNCalculator(model_path, device=self.device, compute_stress=False)
+
+    def encode(self, smiles_list, return_tensor=False, batch_size=32):
+        results = []
+
+        # make batch-wise processing with progress bar
+        for i in tqdm(range(0, len(smiles_list), batch_size), desc="Batch Encoding"):
+            batch = smiles_list[i:i+batch_size]
+            atoms_batch = []
+
+            for smiles in batch:
+                try:
+                    atoms = smiles_to_atoms(smiles)
+                    atoms.calc = self.calculator
+                    atoms_batch.append(atoms)
+                except Exception as e:
+                    print(f"Skipping {smiles}: {e}")
+
+            if atoms_batch:
+                embeddings = [a.get_invariant_embeddings().mean(dim=0).cpu() for a in atoms_batch]
+                batch_tensor = torch.stack(embeddings)
+                results.append(batch_tensor)
+
+        if not results:
+            raise RuntimeError("No valid SMILES could be processed.")
+
+        all_embeddings = torch.cat(results, dim=0)
+        return all_embeddings if return_tensor else pd.DataFrame(all_embeddings.numpy())
+

models/pos_egnn/posegnn/__init__.py

@@ -0,0 +1,3 @@
+from . import calculator, encoder, model, ops, utils
+
+__all__ = ["calculator", "encoder", "model", "ops", "utils"]

models/pos_egnn/posegnn/calculator.py

@@ -0,0 +1,65 @@
+import numpy as np
+import torch
+from ase import Atoms
+from ase.calculators.calculator import Calculator, all_changes
+from ase.data import atomic_numbers
+from ase.stress import full_3x3_to_voigt_6_stress
+from torch_geometric.data.data import Data
+
+from .model import PosEGNN
+
+
+class PosEGNNCalculator(Calculator):
+    def __init__(self, checkpoint: str, device: str, compute_stress: bool = True, **kwargs):
+        Calculator.__init__(self, **kwargs)
+
+        checkpoint_dict = torch.load(checkpoint, weights_only=True, map_location=device)
+
+        self.model = PosEGNN(checkpoint_dict["config"])
+        self.model.load_state_dict(checkpoint_dict["state_dict"], strict=True)
+        self.model.eval()
+
+        self.model.to(device)
+        self.model.eval()
+
+        self.implemented_properties = ["energy", "forces"]
+        self.implemented_properties += ["stress"] if compute_stress else []
+        self.device = device
+        self.compute_stress = compute_stress
+
+    def calculate(self, atoms=None, properties=None, system_changes=all_changes):
+        Calculator.calculate(self, atoms)
+        self.results = {}
+        data = self._build_data(atoms)
+        out = self.model.compute_properties(data, compute_stress=self.compute_stress)
+
+        # Decoder Forward
+        self.results = {
+            "energy": out["total_energy"].cpu().detach().numpy(),
+            "forces": out["force"].cpu().detach().numpy()
+        }
+        if self.compute_stress:
+            self.results.update({
+                "stress": full_3x3_to_voigt_6_stress(out["stress"].cpu().detach().numpy())
+                })
+
+    def _build_data(self, atoms):
+        z = torch.tensor(np.array([atomic_numbers[symbol] for symbol in atoms.symbols]), device=self.device)
+        box = torch.tensor(atoms.get_cell().tolist(), device=self.device).unsqueeze(0).float()
+        pos = torch.tensor(atoms.get_positions().tolist(), device=self.device).float()
+        batch = torch.zeros(len(z), device=self.device).long()
+        ptr = torch.zeros(1, device=self.device).long()
+        return Data(z=z, pos=pos, box=box, batch=batch, num_graphs=1, ptr=ptr)
+
+
+def get_invariant_embeddings(self):
+    if self.calc is None:
+        raise RuntimeError("No calculator is set.")
+    else:
+        data = self.calc._build_data(self)
+        with torch.no_grad():
+            embeddings = self.calc.model(data)["embedding_0"][..., 1].squeeze(2)
+        return embeddings
+
+
+Atoms.get_invariant_embeddings = get_invariant_embeddings

models/pos_egnn/posegnn/encoder.py

@@ -0,0 +1,621 @@
+"""
+This code was adapted from https://github.com/sarpaykent/GotenNet
+Copyright (c) 2025 Sarp Aykent
+MIT License
+
+GotenNet: Rethinking Efficient 3D Equivariant Graph Neural Networks
+Sarp Aykent and Tian Xia
+https://openreview.net/pdf?id=5wxCQDtbMo
+"""
+
+from functools import partial
+from typing import Callable, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch_geometric.nn import MessagePassing
+from torch_geometric.typing import OptTensor
+from torch_geometric.utils import scatter, softmax
+
+from .ops import (
+    MLP,
+    CosineCutoff,
+    Dense,
+    EdgeInit,
+    NodeInit,
+    TensorInit,
+    TensorLayerNorm,
+    get_weight_init_by_string,
+    parse_update_info,
+    str2act,
+    str2basis,
+)
+
+
+def lmax_tensor_size(lmax):
+    return ((lmax + 1) ** 2) - 1
+
+
+def split_degree(tensor, lmax, dim=-1):  # default to last dim
+    cumsum = 0
+    tensors = []
+    for i in range(1, lmax + 1):
+        count = lmax_tensor_size(i) - lmax_tensor_size(i - 1)
+        # Create slice object for the specified dimension
+        slc = [slice(None)] * tensor.ndim  # Create list of slice(None) for all dims
+        slc[dim] = slice(cumsum, cumsum + count)  # Replace desired dim with actual slice
+        tensors.append(tensor[tuple(slc)])
+        cumsum += count
+    return tensors
+
+
+class GATA(MessagePassing):
+    def __init__(
+        self,
+        n_atom_basis: int,
+        activation: Callable,
+        weight_init=nn.init.xavier_uniform_,
+        bias_init=nn.init.zeros_,
+        aggr="add",
+        node_dim=0,
+        epsilon: float = 1e-7,
+        layer_norm=False,
+        vector_norm=False,
+        cutoff=5.0,
+        num_heads=8,
+        dropout=0.0,
+        edge_updates=True,
+        last_layer=False,
+        scale_edge=True,
+        edge_ln="",
+        evec_dim=None,
+        emlp_dim=None,
+        sep_vecj=True,
+        lmax=1,
+    ):
+        """
+        Args:
+            n_atom_basis (int): Number of features to describe atomic environments.
+            activation (Callable): Activation function to be used. If None, no activation function is used.
+            weight_init (Callable): Weight initialization function.
+            bias_init (Callable): Bias initialization function.
+            aggr (str): Aggregation method ('add', 'mean' or 'max').
+            node_dim (int): The axis along which to aggregate.
+        """
+        super(GATA, self).__init__(aggr=aggr, node_dim=node_dim)
+        self.lmax = lmax
+        self.sep_vecj = sep_vecj
+        self.epsilon = epsilon
+        self.last_layer = last_layer
+        self.edge_updates = edge_updates
+        self.scale_edge = scale_edge
+        self.activation = activation
+
+        self.update_info = parse_update_info(edge_updates)
+
+        self.dropout = dropout
+        self.n_atom_basis = n_atom_basis
+
+        InitDense = partial(Dense, weight_init=weight_init, bias_init=bias_init)
+        self.gamma_s = nn.Sequential(
+            InitDense(n_atom_basis, n_atom_basis, activation=activation),
+            InitDense(n_atom_basis, 3 * n_atom_basis, activation=None),
+        )
+
+        self.num_heads = num_heads
+        self.q_w = InitDense(n_atom_basis, n_atom_basis, activation=None)
+        self.k_w = InitDense(n_atom_basis, n_atom_basis, activation=None)
+
+        self.gamma_v = nn.Sequential(
+            InitDense(n_atom_basis, n_atom_basis, activation=activation),
+            InitDense(n_atom_basis, 3 * n_atom_basis, activation=None),
+        )
+
+        self.phik_w_ra = InitDense(
+            n_atom_basis,
+            n_atom_basis,
+            activation=activation,
+        )
+
+        InitMLP = partial(MLP, weight_init=weight_init, bias_init=bias_init)
+
+        self.edge_vec_dim = n_atom_basis if evec_dim is None else evec_dim
+        self.edge_mlp_dim = n_atom_basis if emlp_dim is None else emlp_dim
+        if not self.last_layer and self.edge_updates:
+            if self.update_info["mlp"] or self.update_info["mlpa"]:
+                dims = [n_atom_basis, self.edge_mlp_dim, n_atom_basis]
+            else:
+                dims = [n_atom_basis, n_atom_basis]
+            self.edge_attr_up = InitMLP(
+                dims, activation=activation, last_activation=None if self.update_info["mlp"] else self.activation, norm=edge_ln
+            )
+            self.vecq_w = InitDense(n_atom_basis, self.edge_vec_dim, activation=None, bias=False)
+
+            if self.sep_vecj:
+                self.veck_w = nn.ModuleList(
+                    [InitDense(n_atom_basis, self.edge_vec_dim, activation=None, bias=False) for i in range(self.lmax)]
+                )
+            else:
+                self.veck_w = InitDense(n_atom_basis, self.edge_vec_dim, activation=None, bias=False)
+
+            if self.update_info["lin_w"] > 0:
+                modules = []
+                if self.update_info["lin_w"] % 10 == 2:
+                    modules.append(self.activation)
+                self.lin_w_linear = InitDense(
+                    self.edge_vec_dim,
+                    n_atom_basis,
+                    activation=None,
+                    norm="layer" if self.update_info["lin_w"] == 2 else "",  # lin_ln in original code but error
+                )
+                modules.append(self.lin_w_linear)
+                self.lin_w = nn.Sequential(*modules)
+
+        self.down_proj = nn.Identity()
+
+        self.cutoff = CosineCutoff(cutoff)
+        self._alpha = None
+
+        self.w_re = InitDense(
+            n_atom_basis,
+            n_atom_basis * 3,
+            None,
+        )
+
+        self.layernorm_ = layer_norm
+        self.vector_norm_ = vector_norm
+
+        if layer_norm:
+            self.layernorm = nn.LayerNorm(n_atom_basis)
+        else:
+            self.layernorm = nn.Identity()
+        if vector_norm:
+            self.tln = TensorLayerNorm(n_atom_basis, trainable=False)
+        else:
+            self.tln = nn.Identity()
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.layernorm_:
+            self.layernorm.reset_parameters()
+        if self.vector_norm_:
+            self.tln.reset_parameters()
+        for l in self.gamma_s:  # noqa: E741
+            l.reset_parameters()
+
+        self.q_w.reset_parameters()
+        self.k_w.reset_parameters()
+        for l in self.gamma_v:  # noqa: E741
+            l.reset_parameters()
+        # self.v_w.reset_parameters()
+        # self.out_w.reset_parameters()
+        self.w_re.reset_parameters()
+
+        if not self.last_layer and self.edge_updates:
+            self.edge_attr_up.reset_parameters()
+            self.vecq_w.reset_parameters()
+
+            if self.sep_vecj:
+                for w in self.veck_w:
+                    w.reset_parameters()
+            else:
+                self.veck_w.reset_parameters()
+
+            if self.update_info["lin_w"] > 0:
+                self.lin_w_linear.reset_parameters()
+
+    def forward(
+        self,
+        edge_index,
+        s: torch.Tensor,
+        t: torch.Tensor,
+        dir_ij: torch.Tensor,
+        r_ij: torch.Tensor,
+        d_ij: torch.Tensor,
+        num_edges_expanded: torch.Tensor,
+    ):
+        """Compute interaction output."""
+        s = self.layernorm(s)
+        t = self.tln(t)
+
+        q = self.q_w(s).reshape(-1, self.num_heads, self.n_atom_basis // self.num_heads)
+        k = self.k_w(s).reshape(-1, self.num_heads, self.n_atom_basis // self.num_heads)
+
+        x = self.gamma_s(s)
+        val = self.gamma_v(s)
+        f_ij = r_ij
+        r_ij_attn = self.phik_w_ra(r_ij)
+        r_ij = self.w_re(r_ij)
+
+        # propagate_type: (x: Tensor, ten: Tensor, q:Tensor, k:Tensor, val:Tensor, r_ij: Tensor, r_ij_attn: Tensor, d_ij:Tensor, dir_ij: Tensor, num_edges_expanded: Tensor)
+        su, tu = self.propagate(
+            edge_index=edge_index,
+            x=x,
+            q=q,
+            k=k,
+            val=val,
+            ten=t,
+            r_ij=r_ij,
+            r_ij_attn=r_ij_attn,
+            d_ij=d_ij,
+            dir_ij=dir_ij,
+            num_edges_expanded=num_edges_expanded,
+        )  # , f_ij=f_ij
+
+        s = s + su
+        t = t + tu
+
+        if not self.last_layer and self.edge_updates:
+            vec = t
+
+            w1 = self.vecq_w(vec)
+            if self.sep_vecj:
+                vec_split = split_degree(vec, self.lmax, dim=1)
+                w_out = torch.concat([w(vec_split[i]) for i, w in enumerate(self.veck_w)], dim=1)
+
+            else:
+                w_out = self.veck_w(vec)
+
+            # edge_updater_type: (w1: Tensor, w2:Tensor,  d_ij: Tensor, f_ij: Tensor)
+            df_ij = self.edge_updater(edge_index, w1=w1, w2=w_out, d_ij=dir_ij, f_ij=f_ij)
+            df_ij = f_ij + df_ij
+            self._alpha = None
+            return s, t, df_ij
+        else:
+            self._alpha = None
+            return s, t, f_ij
+
+        # return s, t
+
+    def message(
+        self,
+        edge_index,
+        x_i: torch.Tensor,
+        x_j: torch.Tensor,
+        q_i: torch.Tensor,
+        k_j: torch.Tensor,
+        val_j: torch.Tensor,
+        ten_j: torch.Tensor,
+        r_ij: torch.Tensor,
+        r_ij_attn: torch.Tensor,
+        d_ij: torch.Tensor,
+        dir_ij: torch.Tensor,
+        num_edges_expanded: torch.Tensor,
+        index: torch.Tensor,
+        ptr: OptTensor,
+        dim_size: Optional[int],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute message passing.
+        """
+
+        r_ij_attn = r_ij_attn.reshape(-1, self.num_heads, self.n_atom_basis // self.num_heads)
+        attn = (q_i * k_j * r_ij_attn).sum(dim=-1, keepdim=True)
+
+        attn = softmax(attn, index, ptr, dim_size)
+
+        # Normalize the attention scores
+        if self.scale_edge:
+            norm = torch.sqrt(num_edges_expanded.reshape(-1, 1, 1)) / np.sqrt(self.n_atom_basis)
+        else:
+            norm = 1.0 / np.sqrt(self.n_atom_basis)
+        attn = attn * norm
+        self._alpha = attn
+        attn = F.dropout(attn, p=self.dropout, training=self.training)
+
+        self_attn = attn * val_j.reshape(-1, self.num_heads, (self.n_atom_basis * 3) // self.num_heads)
+        SEA = self_attn.reshape(-1, 1, self.n_atom_basis * 3)
+
+        x = SEA + (r_ij.unsqueeze(1) * x_j * self.cutoff(d_ij.unsqueeze(-1).unsqueeze(-1)))
+
+        o_s, o_d, o_t = torch.split(x, self.n_atom_basis, dim=-1)
+        dmu = o_d * dir_ij[..., None] + o_t * ten_j
+        return o_s, dmu
+
+    @staticmethod
+    def rej(vec, d_ij):
+        vec_proj = (vec * d_ij.unsqueeze(2)).sum(dim=1, keepdim=True)
+        return vec - vec_proj * d_ij.unsqueeze(2)
+
+    def edge_update(self, w1_i, w2_j, w3_j, d_ij, f_ij):
+        if self.sep_vecj:
+            vi = w1_i
+            vj = w2_j
+            vi_split = split_degree(vi, self.lmax, dim=1)
+            vj_split = split_degree(vj, self.lmax, dim=1)
+            d_ij_split = split_degree(d_ij, self.lmax, dim=1)
+
+            pairs = []
+            for i in range(len(vi_split)):
+                if self.update_info["rej"]:
+                    w1 = self.rej(vi_split[i], d_ij_split[i])
+                    w2 = self.rej(vj_split[i], -d_ij_split[i])
+                    pairs.append((w1, w2))
+                else:
+                    w1 = vi_split[i]
+                    w2 = vj_split[i]
+                    pairs.append((w1, w2))
+        elif not self.update_info["rej"]:
+            w1 = w1_i
+            w2 = w2_j
+            pairs = [(w1, w2)]
+        else:
+            w1 = self.rej(w1_i, d_ij)
+            w2 = self.rej(w2_j, -d_ij)
+            pairs = [(w1, w2)]
+
+        w_dot_sum = None
+        for el in pairs:
+            w1, w2 = el
+            w_dot = (w1 * w2).sum(dim=1)
+            if w_dot_sum is None:
+                w_dot_sum = w_dot
+            else:
+                w_dot_sum = w_dot_sum + w_dot
+        w_dot = w_dot_sum
+        if self.update_info["lin_w"] > 0:
+            w_dot = self.lin_w(w_dot)
+
+        if self.update_info["gated"] == "gatedt":
+            w_dot = torch.tanh(w_dot)
+        elif self.update_info["gated"] == "gated":
+            w_dot = torch.sigmoid(w_dot)
+        elif self.update_info["gated"] == "act":
+            w_dot = self.activation(w_dot)
+
+        df_ij = self.edge_attr_up(f_ij) * w_dot
+        return df_ij
+
+    # noinspection PyMethodOverriding
+    def aggregate(
+        self,
+        features: Tuple[torch.Tensor, torch.Tensor],
+        index: torch.Tensor,
+        ptr: Optional[torch.Tensor],
+        dim_size: Optional[int],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        x, vec = features
+        x = scatter(x, index, dim=self.node_dim, dim_size=dim_size, reduce=self.aggr)
+        vec = scatter(vec, index, dim=self.node_dim, dim_size=dim_size, reduce=self.aggr)
+        return x, vec
+
+    def update(self, inputs: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+        return inputs
+
+
+class EQFF(nn.Module):
+    def __init__(
+        self,
+        n_atom_basis: int,
+        activation: Callable,
+        epsilon: float = 1e-8,
+        weight_init=nn.init.xavier_uniform_,
+        bias_init=nn.init.zeros_,
+        vec_dim=None,
+    ):
+        """Equiavariant Feed Forward layer."""
+        super(EQFF, self).__init__()
+        self.n_atom_basis = n_atom_basis
+
+        InitDense = partial(Dense, weight_init=weight_init, bias_init=bias_init)
+
+        vec_dim = n_atom_basis if vec_dim is None else vec_dim
+        context_dim = 2 * n_atom_basis
+
+        self.gamma_m = nn.Sequential(
+            InitDense(context_dim, n_atom_basis, activation=activation),
+            InitDense(n_atom_basis, 2 * n_atom_basis, activation=None),
+        )
+        self.w_vu = InitDense(n_atom_basis, vec_dim, activation=None, bias=False)
+
+        self.epsilon = epsilon
+
+    def reset_parameters(self):
+        self.w_vu.reset_parameters()
+        for l in self.gamma_m:  # noqa: E741
+            l.reset_parameters()
+
+    def forward(self, s, v):
+        """Compute Equivariant Feed Forward output."""
+
+        t_prime = self.w_vu(v)
+        t_prime_mag = torch.sqrt(torch.sum(t_prime**2, dim=-2, keepdim=True) + self.epsilon)
+        combined = [s, t_prime_mag]
+        combined_tensor = torch.cat(combined, dim=-1)
+        m12 = self.gamma_m(combined_tensor)
+
+        m_1, m_2 = torch.split(m12, self.n_atom_basis, dim=-1)
+        delta_v = m_2 * t_prime
+
+        s = s + m_1
+        v = v + delta_v
+
+        return s, v
+
+
+class GotenNet(nn.Module):
+    def __init__(
+        self,
+        hidden_channels: int = 128,
+        num_layers: int = 8, 
+        radial_basis: Union[Callable, str] = "BesselBasis",
+        n_rbf: int = 20,
+        cutoff: float = 5.0,
+        activation: Optional[Union[Callable, str]] = F.silu,
+        max_z: int = 100,
+        epsilon: float = 1e-8,
+        weight_init=nn.init.xavier_uniform_,
+        bias_init=nn.init.zeros_,
+        int_layer_norm=False,
+        int_vector_norm=False,
+        before_mixing_layer_norm=False,
+        after_mixing_layer_norm=False,
+        num_heads=8,
+        attn_dropout=0.0,
+        edge_updates=True,
+        scale_edge=True,
+        lmax=2,
+        aggr="add",
+        edge_ln="",
+        evec_dim=None,
+        emlp_dim=None,
+        sep_int_vec=True,
+    ):
+        """
+        Representation for GotenNet
+        """
+        super(GotenNet, self).__init__()
+
+        self.scale_edge = scale_edge
+        if type(weight_init) == str:  # noqa: E721
+            # print(f"Using {weight_init} weight initialization")
+            weight_init = get_weight_init_by_string(weight_init)
+
+        if type(bias_init) == str:  # noqa: E721
+            bias_init = get_weight_init_by_string(bias_init)
+
+        if type(activation) is str:
+            activation = str2act(activation)
+
+        self.n_atom_basis = self.hidden_dim = hidden_channels
+        self.n_interactions = num_layers
+        self.cutoff = cutoff
+
+        self.neighbor_embedding = NodeInit(
+            [self.hidden_dim // 2, self.hidden_dim],
+            n_rbf,
+            self.cutoff,
+            max_z=max_z,
+            weight_init=weight_init,
+            bias_init=bias_init,
+            concat=False,
+            proj_ln="layer",
+            activation=activation,
+        )
+        self.edge_embedding = EdgeInit(
+            n_rbf, [self.hidden_dim // 2, self.hidden_dim], weight_init=weight_init, bias_init=bias_init, proj_ln=""
+        )
+
+        radial_basis = str2basis(radial_basis)
+        self.radial_basis = radial_basis(cutoff=self.cutoff, n_rbf=n_rbf)
+
+        self.embedding = nn.Embedding(max_z, self.n_atom_basis, padding_idx=0)
+
+        self.tensor_init = TensorInit(l=lmax)
+
+        self.gata = nn.ModuleList(
+            [
+                GATA(
+                    n_atom_basis=self.n_atom_basis,
+                    activation=activation,
+                    aggr=aggr,
+                    weight_init=weight_init,
+                    bias_init=bias_init,
+                    layer_norm=int_layer_norm,
+                    vector_norm=int_vector_norm,
+                    cutoff=self.cutoff,
+                    epsilon=epsilon,
+                    num_heads=num_heads,
+                    dropout=attn_dropout,
+                    edge_updates=edge_updates,
+                    last_layer=(i == self.n_interactions - 1),
+                    scale_edge=scale_edge,
+                    edge_ln=edge_ln,
+                    evec_dim=evec_dim,
+                    emlp_dim=emlp_dim,
+                    sep_vecj=sep_int_vec,
+                    lmax=lmax,
+                )
+                for i in range(self.n_interactions)
+            ]
+        )
+
+        self.eqff = nn.ModuleList(
+            [
+                EQFF(n_atom_basis=self.n_atom_basis, activation=activation, epsilon=epsilon, weight_init=weight_init, bias_init=bias_init)
+                for i in range(self.n_interactions)
+            ]
+        )
+
+        # Extra layer norms for the scalar quantities
+        if before_mixing_layer_norm:
+            self.before_mixing_ln = nn.LayerNorm(self.n_atom_basis)
+        else:
+            self.before_mixing_ln = nn.Identity()
+
+        if after_mixing_layer_norm:
+            self.after_mixing_ln = nn.LayerNorm(self.n_atom_basis)
+        else:
+            self.after_mixing_ln = nn.Identity()
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.edge_embedding.reset_parameters()
+        self.neighbor_embedding.reset_parameters()
+        for l in self.gata:  # noqa: E741
+            l.reset_parameters()
+        for l in self.eqff:  # noqa: E741
+            l.reset_parameters()
+
+        if not isinstance(self.before_mixing_ln, nn.Identity):
+            self.before_mixing_ln.reset_parameters()
+        if not isinstance(self.after_mixing_ln, nn.Identity):
+            self.after_mixing_ln.reset_parameters()
+
+    def forward(self, z, pos, cutoff_edge_index, cutoff_edge_distance, cutoff_edge_vec):
+        q = self.embedding(z)[:]
+
+        edge_attr = self.radial_basis(cutoff_edge_distance)
+
+        q = self.neighbor_embedding(z, q, cutoff_edge_index, cutoff_edge_distance, edge_attr)
+        edge_attr = self.edge_embedding(cutoff_edge_index, edge_attr, q)
+        mask = cutoff_edge_index[0] != cutoff_edge_index[1]
+        # direction vector
+        dist = torch.norm(cutoff_edge_vec[mask], dim=1).unsqueeze(1)
+        cutoff_edge_vec[mask] = cutoff_edge_vec[mask] / dist
+
+        cutoff_edge_vec = self.tensor_init(cutoff_edge_vec)
+        equi_dim = ((self.tensor_init.l + 1) ** 2) - 1
+        # count number of edges for each node
+        num_edges = scatter(torch.ones_like(cutoff_edge_distance), cutoff_edge_index[0], dim=0, reduce="sum")
+        # the shape of num edges is [num_nodes, 1], we want to expand this to [num_edges, 1]
+        # Map num_edges back to the shape of attn using cutoff_edge_index
+        num_edges_expanded = num_edges[cutoff_edge_index[0]]
+
+        qs = q.shape
+        mu = torch.zeros((qs[0], equi_dim, qs[1]), device=q.device)
+        q.unsqueeze_(1)
+
+        layer_outputs = []
+
+        for i, (interaction, mixing) in enumerate(zip(self.gata, self.eqff)):
+            q, mu, edge_attr = interaction(
+                cutoff_edge_index,
+                q,
+                mu,
+                dir_ij=cutoff_edge_vec,
+                r_ij=edge_attr,
+                d_ij=cutoff_edge_distance,
+                num_edges_expanded=num_edges_expanded,
+            )
+
+            q = self.before_mixing_ln(q)
+            q, mu = mixing(q, mu)
+            q = self.after_mixing_ln(q)
+
+            # Collect all scalars for inter-layer read-outs
+            layer_outputs.append(q.squeeze(1))
+
+        # q = q.squeeze(1)
+
+        layer_outputs = torch.stack(layer_outputs, dim=-1)
+
+        output_dict = {}
+        output_dict["embedding_0"] = layer_outputs.unsqueeze(2)  # [n_nodes, n_features, dimension of irrep, n_layers]
+        # This is a scalar so a single irrep
+
+        return output_dict

models/pos_egnn/posegnn/model.py

@@ -0,0 +1,99 @@
+from torch import nn
+import torch
+from .encoder import GotenNet
+from .utils import get_symmetric_displacement, BatchedPeriodicDistance, ACT_CLASS_MAPPING
+#from torch_scatter import scatter
+
+class NodeInvariantReadout(nn.Module):
+    def __init__(self, in_channels, num_residues, hidden_channels, out_channels, activation):
+        super().__init__()
+
+        self.linears = nn.ModuleList([nn.Linear(in_channels, out_channels) for _ in range(num_residues - 1)])
+
+        # Define the nonlinear layer for the last layer's output
+        self.non_linear = nn.Sequential(
+            nn.Linear(in_channels, hidden_channels),
+            ACT_CLASS_MAPPING[activation](),
+            nn.Linear(hidden_channels, out_channels),
+        )
+
+    def forward(self, embedding_0):
+        layer_outputs = embedding_0.squeeze(2)  # [n_nodes, in_channels, num_residues]
+
+        processed_outputs = []
+        for i, linear in enumerate(self.linears):
+            processed_outputs.append(linear(layer_outputs[:, :, i]))
+
+        processed_outputs.append(self.non_linear(layer_outputs[:, :, -1]))
+        output = torch.stack(processed_outputs, dim=0).sum(dim=0).squeeze(-1)
+        return output
+
+class PosEGNN(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.distance = BatchedPeriodicDistance(config["encoder"]["cutoff"])
+        self.encoder = GotenNet(**config["encoder"])
+        self.readout = NodeInvariantReadout(**config["decoder"])
+        self.register_buffer("e0_mean", torch.tensor(config["e0_mean"]))
+        self.register_buffer("atomic_res_total_mean", torch.tensor(config["atomic_res_total_mean"]))
+        self.register_buffer("atomic_res_total_std", torch.tensor(config["atomic_res_total_std"]))
+
+    def forward(self, data):
+        data.pos.requires_grad_(True)
+
+        data.pos, data.box, data.displacements = get_symmetric_displacement(data.pos, data.box, data.num_graphs, data.batch)
+
+        data.cutoff_edge_index, data.cutoff_edge_distance, data.cutoff_edge_vec, data.cutoff_shifts_idx = self.distance(
+            data.pos, data.box, data.batch
+        )
+
+        embedding_dict = self.encoder(data.z, data.pos, data.cutoff_edge_index, data.cutoff_edge_distance, data.cutoff_edge_vec)
+
+        return embedding_dict
+
+    def compute_properties(self, data, compute_stress = True):
+        output = {}
+        
+        embedding_dict = self.forward(data)
+        embedding_0 = embedding_dict["embedding_0"]
+
+        # Compute energy
+        node_e_res = self.readout(embedding_0)
+
+        node_e_res = node_e_res * self.atomic_res_total_std + self.atomic_res_total_mean
+        total_e_res = scatter(src=node_e_res, index=data["batch"], dim=0, reduce="sum")
+
+        node_e0 = self.e0_mean[data.z]
+        total_e0 = scatter(src=node_e0, index=data["batch"], dim=0, reduce="sum")
+
+        total_energy = total_e0 + total_e_res
+        output["total_energy"] = total_energy
+
+        # Compute gradients
+        if compute_stress:
+            inputs = [data.pos, data.displacements]
+            compute_stress = True
+        else:
+            inputs = [data.pos]
+
+        grad_outputs = torch.autograd.grad(
+            outputs=[total_energy],
+            inputs=inputs,
+            grad_outputs=[torch.ones_like(total_energy)],
+            retain_graph=self.training,
+            create_graph=self.training,
+        )
+
+        # Get forces and stresses
+        if compute_stress:
+            force, virial = grad_outputs
+            stress = virial / torch.det(data.box).abs().view(-1, 1, 1)
+            stress = torch.where(torch.abs(stress) < 1e10, stress, torch.zeros_like(stress))
+            output["force"] = -force
+            output["stress"] = -stress
+        else:
+            force = grad_outputs[0]
+            output["force"] = -force        
+        
+        return output

models/pos_egnn/posegnn/ops.py

@@ -0,0 +1,1584 @@
+"""
+This code was adapted from https://github.com/sarpaykent/GotenNet
+Copyright (c) 2025 Sarp Aykent
+MIT License
+
+GotenNet: Rethinking Efficient 3D Equivariant Graph Neural Networks
+Sarp Aykent and Tian Xia
+https://openreview.net/pdf?id=5wxCQDtbMo
+"""
+
+from __future__ import absolute_import, division, print_function
+
+import inspect
+import math
+from functools import partial
+from typing import List
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch import nn as nn
+from torch.nn.init import constant_, xavier_uniform_
+from torch_geometric.nn import MessagePassing
+from torch_geometric.nn.inits import glorot_orthogonal
+from torch_geometric.nn.models.schnet import ShiftedSoftplus
+#from torch_scatter import scatter
+
+zeros_initializer = partial(constant_, val=0.0)
+
+
+def centralize(
+    batch,
+    key: str,
+    batch_index: torch.Tensor,
+):  # note: cannot make assumptions on output shape
+    # derive centroid of each batch element, and center entities using corresponding centroids
+    entities_centroid = scatter(batch[key], batch_index, dim=0, reduce="mean")  # e.g., [batch_size, 3]
+    entities_centered = batch[key] - entities_centroid[batch_index]
+
+    return entities_centroid, entities_centered
+
+
+def decentralize(
+    positions: torch.Tensor,
+    batch_index: torch.Tensor,
+    entities_centroid: torch.Tensor,
+) -> torch.Tensor:  # note: cannot make assumptions on output shape
+    entities_centered = positions + entities_centroid[batch_index]
+    return entities_centered
+
+
+def parse_update_info(edge_updates):
+    update_info = {
+        "gated": False,
+        "rej": True,
+        "vec_norm": False,
+        "mlp": False,
+        "mlpa": False,
+        "lin_w": 0,
+        "drej": False,
+    }
+    if isinstance(edge_updates, str):
+        update_parts = edge_updates.split("_")
+    else:
+        update_parts = []
+
+    allowed_parts = ["gated", "gatedt", "norej", "mlp", "mlpa", "act", "linw", "linwa", "drej"]
+    if not all([part in allowed_parts for part in update_parts]):
+        raise ValueError(f"Invalid edge update parts. Allowed parts are {allowed_parts}")
+
+    if "gated" in update_parts:
+        update_info["gated"] = "gated"
+    if "gatedt" in update_parts:
+        update_info["gated"] = "gatedt"
+    if "act" in update_parts:
+        update_info["gated"] = "act"
+    if "norej" in update_parts:
+        update_info["rej"] = False
+    if "mlp" in update_parts:
+        update_info["mlp"] = True
+    if "mlpa" in update_parts:
+        update_info["mlpa"] = True
+    if "linw" in update_parts:
+        update_info["lin_w"] = 1
+    if "linwa" in update_parts:
+        update_info["lin_w"] = 2
+    if "drej" in update_parts:
+        update_info["drej"] = True
+    return update_info
+
+
+class SmoothLeakyReLU(torch.nn.Module):
+    def __init__(self, negative_slope=0.2):
+        super().__init__()
+        self.alpha = negative_slope
+
+    def forward(self, x):
+        x1 = ((1 + self.alpha) / 2) * x
+        x2 = ((1 - self.alpha) / 2) * x * (2 * torch.sigmoid(x) - 1)
+        return x1 + x2
+
+    def extra_repr(self):
+        return "negative_slope={}".format(self.alpha)
+
+
+def shifted_softplus(x: torch.Tensor):
+    return F.softplus(x) - math.log(2.0)
+
+
+class PolynomialCutoff(nn.Module):
+    def __init__(self, cutoff, p: int = 6):
+        super(PolynomialCutoff, self).__init__()
+        self.cutoff = cutoff
+        self.p = p
+
+    @staticmethod
+    def polynomial_cutoff(r: Tensor, rcut: float, p: float = 6.0) -> Tensor:
+        """
+        Polynomial cutoff, as proposed in DimeNet: https://arxiv.org/abs/2003.03123
+        """
+        if not p >= 2.0:
+            # replace below with logger error
+            print(f"Exponent p={p} has to be >= 2.")
+            print("Exiting code.")
+
+            print(f"Exponent p={p} has to be >= 2.")
+            print("Exiting code.")
+            exit()
+
+        rscaled = r / rcut
+
+        out = 1.0
+        out = out - (((p + 1.0) * (p + 2.0) / 2.0) * torch.pow(rscaled, p))
+        out = out + (p * (p + 2.0) * torch.pow(rscaled, p + 1.0))
+        out = out - ((p * (p + 1.0) / 2) * torch.pow(rscaled, p + 2.0))
+
+        return out * (rscaled < 1.0).float()
+
+    def forward(self, r):
+        return self.polynomial_cutoff(r=r, rcut=self.cutoff, p=self.p)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(cutoff={self.cutoff}, p={self.p})"
+
+
+class CosineCutoff(nn.Module):
+    def __init__(self, cutoff):
+        super(CosineCutoff, self).__init__()
+
+        if isinstance(cutoff, torch.Tensor):
+            cutoff = cutoff.item()
+        self.cutoff = cutoff
+
+    def forward(self, distances):
+        cutoffs = 0.5 * (torch.cos(distances * math.pi / self.cutoff) + 1.0)
+        cutoffs = cutoffs * (distances < self.cutoff).float()
+        return cutoffs
+
+
+class ScaleShift(nn.Module):
+    r"""Scale and shift layer for standardization.
+
+    .. math::
+       y = x \times \sigma + \mu
+
+    Args:
+        mean (torch.Tensor): mean value :math:`\mu`.
+        stddev (torch.Tensor): standard deviation value :math:`\sigma`.
+
+    """
+
+    def __init__(self, mean, stddev):
+        super(ScaleShift, self).__init__()
+        if isinstance(mean, float):
+            mean = torch.FloatTensor([mean])
+        if isinstance(stddev, float):
+            stddev = torch.FloatTensor([stddev])
+        self.register_buffer("mean", mean)
+        self.register_buffer("stddev", stddev)
+
+    def forward(self, input):
+        """Compute layer output.
+
+        Args:
+            input (torch.Tensor): input data.
+
+        Returns:
+            torch.Tensor: layer output.
+
+        """
+        y = input * self.stddev + self.mean
+        return y
+
+
+class GetItem(nn.Module):
+    """Extraction layer to get an item from SchNetPack dictionary of input tensors.
+    Args:
+        key (str): Property to be extracted from SchNetPack input tensors.
+    """
+
+    def __init__(self, key):
+        super(GetItem, self).__init__()
+        self.key = key
+
+    def forward(self, inputs):
+        """Compute layer output.
+        Args:
+            inputs (dict of torch.Tensor): SchNetPack dictionary of input tensors.
+        Returns:
+            torch.Tensor: layer output.
+        """
+        return inputs[self.key]
+
+
+class SchnetMLP(nn.Module):
+    """Multiple layer fully connected perceptron neural network.
+    Args:
+        n_in (int): number of input nodes.
+        n_out (int): number of output nodes.
+        n_hidden (list of int or int, optional): number hidden layer nodes.
+            If an integer, same number of node is used for all hidden layers resulting
+            in a rectangular network.
+            If None, the number of neurons is divided by two after each layer starting
+            n_in resulting in a pyramidal network.
+        n_layers (int, optional): number of layers.
+        activation (callable, optional): activation function. All hidden layers would
+            the same activation function except the output layer that does not apply
+            any activation function.
+    """
+
+    def __init__(self, n_in, n_out, n_hidden=None, n_layers=2, activation=shifted_softplus):
+        super(SchnetMLP, self).__init__()
+        # get list of number of nodes in input, hidden & output layers
+        if n_hidden is None:
+            c_neurons = n_in
+            self.n_neurons = []
+            for i in range(n_layers):
+                self.n_neurons.append(c_neurons)
+                c_neurons = c_neurons // 2
+            self.n_neurons.append(n_out)
+        else:
+            # get list of number of nodes hidden layers
+            if type(n_hidden) is int:
+                n_hidden = [n_hidden] * (n_layers - 1)
+            self.n_neurons = [n_in] + n_hidden + [n_out]
+
+        # assign a Dense layer (with activation function) to each hidden layer
+        layers = [Dense(self.n_neurons[i], self.n_neurons[i + 1], activation=activation) for i in range(n_layers - 1)]
+        # assign a Dense layer (without activation function) to the output layer
+        layers.append(Dense(self.n_neurons[-2], self.n_neurons[-1], activation=None))
+        # put all layers together to make the network
+        self.out_net = nn.Sequential(*layers)
+
+    def forward(self, inputs):
+        """Compute neural network output.
+        Args:
+            inputs (torch.Tensor): network input.
+        Returns:
+            torch.Tensor: network output.
+        """
+        return self.out_net(inputs)
+
+
+def scaled_silu(x, scale=0.6):
+    return F.silu(x) * scale
+
+
+def gaussian_rbf(inputs: torch.Tensor, offsets: torch.Tensor, widths: torch.Tensor):
+    coeff = -0.5 / torch.pow(widths, 2)
+    diff = inputs[..., None] - offsets
+    y = torch.exp(coeff * torch.pow(diff, 2))
+    return y
+
+
+class GaussianRBF(nn.Module):
+    r"""Gaussian radial basis functions."""
+
+    def __init__(self, n_rbf: int, cutoff: float, start: float = 0.0, trainable: bool = False):
+        """
+        Args:
+            n_rbf: total number of Gaussian functions, :math:`N_g`.
+            cutoff: center of last Gaussian function, :math:`\mu_{N_g}`
+            start: center of first Gaussian function, :math:`\mu_0`.
+            trainable: If True, widths and offset of Gaussian functions
+                are adjusted during training process.
+        """
+        super(GaussianRBF, self).__init__()
+        self.n_rbf = n_rbf
+
+        # compute offset and width of Gaussian functions
+        offset = torch.linspace(start, cutoff, n_rbf)
+        widths = torch.FloatTensor(torch.abs(offset[1] - offset[0]) * torch.ones_like(offset))
+        if trainable:
+            self.widths = nn.Parameter(widths)
+            self.offsets = nn.Parameter(offset)
+        else:
+            self.register_buffer("widths", widths)
+            self.register_buffer("offsets", offset)
+
+    def forward(self, inputs: torch.Tensor):
+        return gaussian_rbf(inputs, self.offsets, self.widths)
+
+
+class BesselBasis(nn.Module):
+    """
+    Sine for radial basis expansion with coulomb decay. (0th order Bessel from DimeNet)
+    """
+
+    def __init__(self, cutoff=5.0, n_rbf=None, trainable=False):
+        """
+        Args:
+            cutoff: radial cutoff
+            n_rbf: number of basis functions.
+        """
+        super(BesselBasis, self).__init__()
+        self.n_rbf = n_rbf
+        # compute offset and width of Gaussian functions
+        freqs = torch.arange(1, n_rbf + 1) * math.pi / cutoff
+        self.register_buffer("freqs", freqs)
+        self.register_buffer("norm1", torch.tensor(1.0))
+
+    def forward(self, inputs):
+        input_size = len(inputs.shape)  # noqa: F841
+        a = self.freqs[None, :]
+        inputs = inputs[..., None]
+        ax = inputs * a
+        sinax = torch.sin(ax)
+
+        norm = torch.where(inputs == 0, self.norm1, inputs)
+        y = sinax / norm
+
+        return y
+
+
+def glorot_orthogonal_wrapper_(tensor, scale=2.0):
+    return glorot_orthogonal(tensor, scale=scale)
+
+
+def _standardize(kernel):
+    """
+    Makes sure that Var(W) = 1 and E[W] = 0
+    """
+    eps = 1e-6
+
+    if len(kernel.shape) == 3:
+        axis = [0, 1]  # last dimension is output dimension
+    else:
+        axis = 1
+
+    var, mean = torch.var_mean(kernel, dim=axis, unbiased=True, keepdim=True)
+    kernel = (kernel - mean) / (var + eps) ** 0.5
+    return kernel
+
+
+def he_orthogonal_init(tensor):
+    """
+    Generate a weight matrix with variance according to He initialization.
+    Based on a random (semi-)orthogonal matrix neural networks
+    are expected to learn better when features are decorrelated
+    (stated by eg. "Reducing overfitting in deep networks by decorrelating representations",
+    "Dropout: a simple way to prevent neural networks from overfitting",
+    "Exact solutions to the nonlinear dynamics of learning in deep linear neural networks")
+    """
+    tensor = torch.nn.init.orthogonal_(tensor)
+
+    if len(tensor.shape) == 3:
+        fan_in = tensor.shape[:-1].numel()
+    else:
+        fan_in = tensor.shape[1]
+
+    with torch.no_grad():
+        tensor.data = _standardize(tensor.data)
+        tensor.data *= (1 / fan_in) ** 0.5
+
+    return tensor
+
+
+def get_weight_init_by_string(init_str):
+    if init_str == "":
+        # Noop
+        return lambda x: x
+    elif init_str == "zeros":
+        return torch.nn.init.zeros_
+    elif init_str == "xavier_uniform":
+        return torch.nn.init.xavier_uniform_
+    elif init_str == "glo_orthogonal":
+        return glorot_orthogonal_wrapper_
+    elif init_str == "he_orthogonal":
+        return he_orthogonal_init
+    else:
+        raise ValueError(f"Unknown initialization {init_str}")
+
+
+class Dense(nn.Linear):
+    r"""Fully connected linear layer with activation function.
+    Barrowed from https://github.com/atomistic-machine-learning/schnetpack/blob/master/src/schnetpack/nn/base.py
+
+    .. math::
+       y = activation(xW^T + b)
+
+    Args:
+        in_features (int): number of input feature :math:`x`.
+        out_features (int): number of output features :math:`y`.
+        bias (bool, optional): if False, the layer will not adapt bias :math:`b`.
+        activation (callable, optional): if None, no activation function is used.
+        weight_init (callable, optional): weight initializer from current weight.
+        bias_init (callable, optional): bias initializer from current bias.
+
+    """
+
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        bias=True,
+        activation=None,
+        weight_init=xavier_uniform_,
+        bias_init=zeros_initializer,
+        norm=None,
+        gain=None,
+    ):
+        # initialize linear layer y = xW^T + b
+        self.weight_init = weight_init
+        self.bias_init = bias_init
+        self.gain = gain
+        super(Dense, self).__init__(in_features, out_features, bias)
+        # Initialize activation function
+        if inspect.isclass(activation):
+            self.activation = activation()
+        self.activation = activation
+
+        if norm == "layer":
+            self.norm = nn.LayerNorm(out_features)
+        elif norm == "batch":
+            self.norm = nn.BatchNorm1d(out_features)
+        elif norm == "instance":
+            self.norm = nn.InstanceNorm1d(out_features)
+        else:
+            self.norm = None
+
+    def reset_parameters(self):
+        """Reinitialize model weight and bias values."""
+        if self.gain:
+            self.weight_init(self.weight, gain=self.gain)
+        else:
+            self.weight_init(self.weight)
+        if self.bias is not None:
+            self.bias_init(self.bias)
+
+    def forward(self, inputs):
+        """Compute layer output.
+
+        Args:
+            inputs (dict of torch.Tensor): batch of input values.
+
+        Returns:
+            torch.Tensor: layer output.
+
+        """
+        # compute linear layer y = xW^T + b
+        y = super(Dense, self).forward(inputs)
+        if self.norm is not None:
+            y = self.norm(y)
+        # add activation function
+        if self.activation:
+            y = self.activation(y)
+        return y
+
+
+class _VDropout(nn.Module):
+    """
+    Vector channel dropout where the elements of each
+    vector channel are dropped together.
+    """
+
+    def __init__(self, drop_rate, scale=True):
+        super(_VDropout, self).__init__()
+        self.drop_rate = drop_rate
+        self.scale = scale
+
+    def forward(self, x, dim=-1):
+        """
+        :param x: `torch.Tensor` corresponding to vector channels
+        """
+        if self.drop_rate == 0:
+            return x
+        device = x.device
+        if not self.training:
+            return x
+
+        shape = list(x.shape)
+        assert shape[dim] == 3, "The dimension must be vector"
+        shape[dim] = 1
+
+        mask = torch.bernoulli((1 - self.drop_rate) * torch.ones(shape, device=device))
+        x = mask * x
+        if self.scale:
+            # scale the output to keep the expected output distribution
+            # same as input distribution. However, this might be harmfuk
+            # for vector space.
+            x = x / (1 - self.drop_rate)
+
+        return x
+
+
+class Dropout(nn.Module):
+    """
+    Combined dropout for tuples (s, V).
+    Takes tuples (s, V) as input and as output.
+    """
+
+    def __init__(self, drop_rate, vector_dropout=True):
+        super(Dropout, self).__init__()
+        self.sdropout = nn.Dropout(drop_rate)
+        if vector_dropout:
+            self.vdropout = _VDropout(drop_rate)
+        else:
+            self.vdropout = lambda x, dim: x
+
+    def forward(self, x):
+        """
+        :param x: tuple (s, V) of `torch.Tensor`,
+                  or single `torch.Tensor`
+                  (will be assumed to be scalar channels)
+        """
+        if type(x) is torch.Tensor:
+            return self.sdropout(x)
+        s, v = x
+        return self.sdropout(s), self.vdropout(v, dim=1)
+
+
+class TensorInit(nn.Module):
+    def __init__(self, l=2):  # noqa: E741
+        super(TensorInit, self).__init__()
+        self.l = l
+
+    def forward(self, edge_vec):
+        edge_sh = self._calculate_components(self.l, edge_vec[..., 0], edge_vec[..., 1], edge_vec[..., 2])
+        return edge_sh
+
+    @property
+    def tensor_size(self):
+        return ((self.l + 1) ** 2) - 1
+
+    @staticmethod
+    def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor:
+        sh_1_0, sh_1_1, sh_1_2 = x, y, z
+
+        if lmax == 1:
+            return torch.stack([sh_1_0, sh_1_1, sh_1_2], dim=-1)
+
+        # (x^2, y^2, z^2) ^2
+
+        sh_2_0 = math.sqrt(3.0) * x * z
+        sh_2_1 = math.sqrt(3.0) * x * y
+        y2 = y.pow(2)
+        x2z2 = x.pow(2) + z.pow(2)
+        sh_2_2 = y2 - 0.5 * x2z2
+        sh_2_3 = math.sqrt(3.0) * y * z
+        sh_2_4 = math.sqrt(3.0) / 2.0 * (z.pow(2) - x.pow(2))
+
+        if lmax == 2:
+            return torch.stack([sh_1_0, sh_1_1, sh_1_2, sh_2_0, sh_2_1, sh_2_2, sh_2_3, sh_2_4], dim=-1)
+
+        # Borrowed from e3nn: https://github.com/e3nn/e3nn/blob/main/e3nn/o3/_spherical_harmonics.py#L188
+        sh_3_0 = (1 / 6) * math.sqrt(42) * (sh_2_0 * z + sh_2_4 * x)
+        sh_3_1 = math.sqrt(7) * sh_2_0 * y
+        sh_3_2 = (1 / 8) * math.sqrt(168) * (4.0 * y2 - x2z2) * x
+        sh_3_3 = (1 / 2) * math.sqrt(7) * y * (2.0 * y2 - 3.0 * x2z2)
+        sh_3_4 = (1 / 8) * math.sqrt(168) * z * (4.0 * y2 - x2z2)
+        sh_3_5 = math.sqrt(7) * sh_2_4 * y
+        sh_3_6 = (1 / 6) * math.sqrt(42) * (sh_2_4 * z - sh_2_0 * x)
+
+        if lmax == 3:
+            return torch.stack(
+                [
+                    sh_1_0,
+                    sh_1_1,
+                    sh_1_2,
+                    sh_2_0,
+                    sh_2_1,
+                    sh_2_2,
+                    sh_2_3,
+                    sh_2_4,
+                    sh_3_0,
+                    sh_3_1,
+                    sh_3_2,
+                    sh_3_3,
+                    sh_3_4,
+                    sh_3_5,
+                    sh_3_6,
+                ],
+                dim=-1,
+            )
+
+        sh_4_0 = (3 / 4) * math.sqrt(2) * (sh_3_0 * z + sh_3_6 * x)
+        sh_4_1 = (3 / 4) * sh_3_0 * y + (3 / 8) * math.sqrt(6) * sh_3_1 * z + (3 / 8) * math.sqrt(6) * sh_3_5 * x
+        sh_4_2 = (
+            -3 / 56 * math.sqrt(14) * sh_3_0 * z
+            + (3 / 14) * math.sqrt(21) * sh_3_1 * y
+            + (3 / 56) * math.sqrt(210) * sh_3_2 * z
+            + (3 / 56) * math.sqrt(210) * sh_3_4 * x
+            + (3 / 56) * math.sqrt(14) * sh_3_6 * x
+        )
+        sh_4_3 = (
+            -3 / 56 * math.sqrt(42) * sh_3_1 * z
+            + (3 / 28) * math.sqrt(105) * sh_3_2 * y
+            + (3 / 28) * math.sqrt(70) * sh_3_3 * x
+            + (3 / 56) * math.sqrt(42) * sh_3_5 * x
+        )
+        sh_4_4 = -3 / 28 * math.sqrt(42) * sh_3_2 * x + (3 / 7) * math.sqrt(7) * sh_3_3 * y - 3 / 28 * math.sqrt(42) * sh_3_4 * z
+        sh_4_5 = (
+            -3 / 56 * math.sqrt(42) * sh_3_1 * x
+            + (3 / 28) * math.sqrt(70) * sh_3_3 * z
+            + (3 / 28) * math.sqrt(105) * sh_3_4 * y
+            - 3 / 56 * math.sqrt(42) * sh_3_5 * z
+        )
+        sh_4_6 = (
+            -3 / 56 * math.sqrt(14) * sh_3_0 * x
+            - 3 / 56 * math.sqrt(210) * sh_3_2 * x
+            + (3 / 56) * math.sqrt(210) * sh_3_4 * z
+            + (3 / 14) * math.sqrt(21) * sh_3_5 * y
+            - 3 / 56 * math.sqrt(14) * sh_3_6 * z
+        )
+        sh_4_7 = -3 / 8 * math.sqrt(6) * sh_3_1 * x + (3 / 8) * math.sqrt(6) * sh_3_5 * z + (3 / 4) * sh_3_6 * y
+        sh_4_8 = (3 / 4) * math.sqrt(2) * (-sh_3_0 * x + sh_3_6 * z)
+        if lmax == 4:
+            return torch.stack(
+                [
+                    sh_1_0,
+                    sh_1_1,
+                    sh_1_2,
+                    sh_2_0,
+                    sh_2_1,
+                    sh_2_2,
+                    sh_2_3,
+                    sh_2_4,
+                    sh_3_0,
+                    sh_3_1,
+                    sh_3_2,
+                    sh_3_3,
+                    sh_3_4,
+                    sh_3_5,
+                    sh_3_6,
+                    sh_4_0,
+                    sh_4_1,
+                    sh_4_2,
+                    sh_4_3,
+                    sh_4_4,
+                    sh_4_5,
+                    sh_4_6,
+                    sh_4_7,
+                    sh_4_8,
+                ],
+                dim=-1,
+            )
+
+        sh_5_0 = (1 / 10) * math.sqrt(110) * (sh_4_0 * z + sh_4_8 * x)
+        sh_5_1 = (1 / 5) * math.sqrt(11) * sh_4_0 * y + (1 / 5) * math.sqrt(22) * sh_4_1 * z + (1 / 5) * math.sqrt(22) * sh_4_7 * x
+        sh_5_2 = (
+            -1 / 30 * math.sqrt(22) * sh_4_0 * z
+            + (4 / 15) * math.sqrt(11) * sh_4_1 * y
+            + (1 / 15) * math.sqrt(154) * sh_4_2 * z
+            + (1 / 15) * math.sqrt(154) * sh_4_6 * x
+            + (1 / 30) * math.sqrt(22) * sh_4_8 * x
+        )
+        sh_5_3 = (
+            -1 / 30 * math.sqrt(66) * sh_4_1 * z
+            + (1 / 15) * math.sqrt(231) * sh_4_2 * y
+            + (1 / 30) * math.sqrt(462) * sh_4_3 * z
+            + (1 / 30) * math.sqrt(462) * sh_4_5 * x
+            + (1 / 30) * math.sqrt(66) * sh_4_7 * x
+        )
+        sh_5_4 = (
+            -1 / 15 * math.sqrt(33) * sh_4_2 * z
+            + (2 / 15) * math.sqrt(66) * sh_4_3 * y
+            + (1 / 15) * math.sqrt(165) * sh_4_4 * x
+            + (1 / 15) * math.sqrt(33) * sh_4_6 * x
+        )
+        sh_5_5 = -1 / 15 * math.sqrt(110) * sh_4_3 * x + (1 / 3) * math.sqrt(11) * sh_4_4 * y - 1 / 15 * math.sqrt(110) * sh_4_5 * z
+        sh_5_6 = (
+            -1 / 15 * math.sqrt(33) * sh_4_2 * x
+            + (1 / 15) * math.sqrt(165) * sh_4_4 * z
+            + (2 / 15) * math.sqrt(66) * sh_4_5 * y
+            - 1 / 15 * math.sqrt(33) * sh_4_6 * z
+        )
+        sh_5_7 = (
+            -1 / 30 * math.sqrt(66) * sh_4_1 * x
+            - 1 / 30 * math.sqrt(462) * sh_4_3 * x
+            + (1 / 30) * math.sqrt(462) * sh_4_5 * z
+            + (1 / 15) * math.sqrt(231) * sh_4_6 * y
+            - 1 / 30 * math.sqrt(66) * sh_4_7 * z
+        )
+        sh_5_8 = (
+            -1 / 30 * math.sqrt(22) * sh_4_0 * x
+            - 1 / 15 * math.sqrt(154) * sh_4_2 * x
+            + (1 / 15) * math.sqrt(154) * sh_4_6 * z
+            + (4 / 15) * math.sqrt(11) * sh_4_7 * y
+            - 1 / 30 * math.sqrt(22) * sh_4_8 * z
+        )
+        sh_5_9 = -1 / 5 * math.sqrt(22) * sh_4_1 * x + (1 / 5) * math.sqrt(22) * sh_4_7 * z + (1 / 5) * math.sqrt(11) * sh_4_8 * y
+        sh_5_10 = (1 / 10) * math.sqrt(110) * (-sh_4_0 * x + sh_4_8 * z)
+        if lmax == 5:
+            return torch.stack(
+                [
+                    sh_1_0,
+                    sh_1_1,
+                    sh_1_2,
+                    sh_2_0,
+                    sh_2_1,
+                    sh_2_2,
+                    sh_2_3,
+                    sh_2_4,
+                    sh_3_0,
+                    sh_3_1,
+                    sh_3_2,
+                    sh_3_3,
+                    sh_3_4,
+                    sh_3_5,
+                    sh_3_6,
+                    sh_4_0,
+                    sh_4_1,
+                    sh_4_2,
+                    sh_4_3,
+                    sh_4_4,
+                    sh_4_5,
+                    sh_4_6,
+                    sh_4_7,
+                    sh_4_8,
+                    sh_5_0,
+                    sh_5_1,
+                    sh_5_2,
+                    sh_5_3,
+                    sh_5_4,
+                    sh_5_5,
+                    sh_5_6,
+                    sh_5_7,
+                    sh_5_8,
+                    sh_5_9,
+                    sh_5_10,
+                ],
+                dim=-1,
+            )
+
+        sh_6_0 = (1 / 6) * math.sqrt(39) * (sh_5_0 * z + sh_5_10 * x)
+        sh_6_1 = (1 / 6) * math.sqrt(13) * sh_5_0 * y + (1 / 12) * math.sqrt(130) * sh_5_1 * z + (1 / 12) * math.sqrt(130) * sh_5_9 * x
+        sh_6_2 = (
+            -1 / 132 * math.sqrt(286) * sh_5_0 * z
+            + (1 / 33) * math.sqrt(715) * sh_5_1 * y
+            + (1 / 132) * math.sqrt(286) * sh_5_10 * x
+            + (1 / 44) * math.sqrt(1430) * sh_5_2 * z
+            + (1 / 44) * math.sqrt(1430) * sh_5_8 * x
+        )
+        sh_6_3 = (
+            -1 / 132 * math.sqrt(858) * sh_5_1 * z
+            + (1 / 22) * math.sqrt(429) * sh_5_2 * y
+            + (1 / 22) * math.sqrt(286) * sh_5_3 * z
+            + (1 / 22) * math.sqrt(286) * sh_5_7 * x
+            + (1 / 132) * math.sqrt(858) * sh_5_9 * x
+        )
+        sh_6_4 = (
+            -1 / 66 * math.sqrt(429) * sh_5_2 * z
+            + (2 / 33) * math.sqrt(286) * sh_5_3 * y
+            + (1 / 66) * math.sqrt(2002) * sh_5_4 * z
+            + (1 / 66) * math.sqrt(2002) * sh_5_6 * x
+            + (1 / 66) * math.sqrt(429) * sh_5_8 * x
+        )
+        sh_6_5 = (
+            -1 / 66 * math.sqrt(715) * sh_5_3 * z
+            + (1 / 66) * math.sqrt(5005) * sh_5_4 * y
+            + (1 / 66) * math.sqrt(3003) * sh_5_5 * x
+            + (1 / 66) * math.sqrt(715) * sh_5_7 * x
+        )
+        sh_6_6 = -1 / 66 * math.sqrt(2145) * sh_5_4 * x + (1 / 11) * math.sqrt(143) * sh_5_5 * y - 1 / 66 * math.sqrt(2145) * sh_5_6 * z
+        sh_6_7 = (
+            -1 / 66 * math.sqrt(715) * sh_5_3 * x
+            + (1 / 66) * math.sqrt(3003) * sh_5_5 * z
+            + (1 / 66) * math.sqrt(5005) * sh_5_6 * y
+            - 1 / 66 * math.sqrt(715) * sh_5_7 * z
+        )
+        sh_6_8 = (
+            -1 / 66 * math.sqrt(429) * sh_5_2 * x
+            - 1 / 66 * math.sqrt(2002) * sh_5_4 * x
+            + (1 / 66) * math.sqrt(2002) * sh_5_6 * z
+            + (2 / 33) * math.sqrt(286) * sh_5_7 * y
+            - 1 / 66 * math.sqrt(429) * sh_5_8 * z
+        )
+        sh_6_9 = (
+            -1 / 132 * math.sqrt(858) * sh_5_1 * x
+            - 1 / 22 * math.sqrt(286) * sh_5_3 * x
+            + (1 / 22) * math.sqrt(286) * sh_5_7 * z
+            + (1 / 22) * math.sqrt(429) * sh_5_8 * y
+            - 1 / 132 * math.sqrt(858) * sh_5_9 * z
+        )
+        sh_6_10 = (
+            -1 / 132 * math.sqrt(286) * sh_5_0 * x
+            - 1 / 132 * math.sqrt(286) * sh_5_10 * z
+            - 1 / 44 * math.sqrt(1430) * sh_5_2 * x
+            + (1 / 44) * math.sqrt(1430) * sh_5_8 * z
+            + (1 / 33) * math.sqrt(715) * sh_5_9 * y
+        )
+        sh_6_11 = -1 / 12 * math.sqrt(130) * sh_5_1 * x + (1 / 6) * math.sqrt(13) * sh_5_10 * y + (1 / 12) * math.sqrt(130) * sh_5_9 * z
+        sh_6_12 = (1 / 6) * math.sqrt(39) * (-sh_5_0 * x + sh_5_10 * z)
+        if lmax == 6:
+            return torch.stack(
+                [
+                    sh_1_0,
+                    sh_1_1,
+                    sh_1_2,
+                    sh_2_0,
+                    sh_2_1,
+                    sh_2_2,
+                    sh_2_3,
+                    sh_2_4,
+                    sh_3_0,
+                    sh_3_1,
+                    sh_3_2,
+                    sh_3_3,
+                    sh_3_4,
+                    sh_3_5,
+                    sh_3_6,
+                    sh_4_0,
+                    sh_4_1,
+                    sh_4_2,
+                    sh_4_3,
+                    sh_4_4,
+                    sh_4_5,
+                    sh_4_6,
+                    sh_4_7,
+                    sh_4_8,
+                    sh_5_0,
+                    sh_5_1,
+                    sh_5_2,
+                    sh_5_3,
+                    sh_5_4,
+                    sh_5_5,
+                    sh_5_6,
+                    sh_5_7,
+                    sh_5_8,
+                    sh_5_9,
+                    sh_5_10,
+                    sh_6_0,
+                    sh_6_1,
+                    sh_6_2,
+                    sh_6_3,
+                    sh_6_4,
+                    sh_6_5,
+                    sh_6_6,
+                    sh_6_7,
+                    sh_6_8,
+                    sh_6_9,
+                    sh_6_10,
+                    sh_6_11,
+                    sh_6_12,
+                ],
+                dim=-1,
+            )
+
+        sh_7_0 = (1 / 14) * math.sqrt(210) * (sh_6_0 * z + sh_6_12 * x)
+        sh_7_1 = (1 / 7) * math.sqrt(15) * sh_6_0 * y + (3 / 7) * math.sqrt(5) * sh_6_1 * z + (3 / 7) * math.sqrt(5) * sh_6_11 * x
+        sh_7_2 = (
+            -1 / 182 * math.sqrt(390) * sh_6_0 * z
+            + (6 / 91) * math.sqrt(130) * sh_6_1 * y
+            + (3 / 91) * math.sqrt(715) * sh_6_10 * x
+            + (1 / 182) * math.sqrt(390) * sh_6_12 * x
+            + (3 / 91) * math.sqrt(715) * sh_6_2 * z
+        )
+        sh_7_3 = (
+            -3 / 182 * math.sqrt(130) * sh_6_1 * z
+            + (3 / 182) * math.sqrt(130) * sh_6_11 * x
+            + (3 / 91) * math.sqrt(715) * sh_6_2 * y
+            + (5 / 182) * math.sqrt(858) * sh_6_3 * z
+            + (5 / 182) * math.sqrt(858) * sh_6_9 * x
+        )
+        sh_7_4 = (
+            (3 / 91) * math.sqrt(65) * sh_6_10 * x
+            - 3 / 91 * math.sqrt(65) * sh_6_2 * z
+            + (10 / 91) * math.sqrt(78) * sh_6_3 * y
+            + (15 / 182) * math.sqrt(78) * sh_6_4 * z
+            + (15 / 182) * math.sqrt(78) * sh_6_8 * x
+        )
+        sh_7_5 = (
+            -5 / 91 * math.sqrt(39) * sh_6_3 * z
+            + (15 / 91) * math.sqrt(39) * sh_6_4 * y
+            + (3 / 91) * math.sqrt(390) * sh_6_5 * z
+            + (3 / 91) * math.sqrt(390) * sh_6_7 * x
+            + (5 / 91) * math.sqrt(39) * sh_6_9 * x
+        )
+        sh_7_6 = (
+            -15 / 182 * math.sqrt(26) * sh_6_4 * z
+            + (12 / 91) * math.sqrt(65) * sh_6_5 * y
+            + (2 / 91) * math.sqrt(1365) * sh_6_6 * x
+            + (15 / 182) * math.sqrt(26) * sh_6_8 * x
+        )
+        sh_7_7 = -3 / 91 * math.sqrt(455) * sh_6_5 * x + (1 / 13) * math.sqrt(195) * sh_6_6 * y - 3 / 91 * math.sqrt(455) * sh_6_7 * z
+        sh_7_8 = (
+            -15 / 182 * math.sqrt(26) * sh_6_4 * x
+            + (2 / 91) * math.sqrt(1365) * sh_6_6 * z
+            + (12 / 91) * math.sqrt(65) * sh_6_7 * y
+            - 15 / 182 * math.sqrt(26) * sh_6_8 * z
+        )
+        sh_7_9 = (
+            -5 / 91 * math.sqrt(39) * sh_6_3 * x
+            - 3 / 91 * math.sqrt(390) * sh_6_5 * x
+            + (3 / 91) * math.sqrt(390) * sh_6_7 * z
+            + (15 / 91) * math.sqrt(39) * sh_6_8 * y
+            - 5 / 91 * math.sqrt(39) * sh_6_9 * z
+        )
+        sh_7_10 = (
+            -3 / 91 * math.sqrt(65) * sh_6_10 * z
+            - 3 / 91 * math.sqrt(65) * sh_6_2 * x
+            - 15 / 182 * math.sqrt(78) * sh_6_4 * x
+            + (15 / 182) * math.sqrt(78) * sh_6_8 * z
+            + (10 / 91) * math.sqrt(78) * sh_6_9 * y
+        )
+        sh_7_11 = (
+            -3 / 182 * math.sqrt(130) * sh_6_1 * x
+            + (3 / 91) * math.sqrt(715) * sh_6_10 * y
+            - 3 / 182 * math.sqrt(130) * sh_6_11 * z
+            - 5 / 182 * math.sqrt(858) * sh_6_3 * x
+            + (5 / 182) * math.sqrt(858) * sh_6_9 * z
+        )
+        sh_7_12 = (
+            -1 / 182 * math.sqrt(390) * sh_6_0 * x
+            + (3 / 91) * math.sqrt(715) * sh_6_10 * z
+            + (6 / 91) * math.sqrt(130) * sh_6_11 * y
+            - 1 / 182 * math.sqrt(390) * sh_6_12 * z
+            - 3 / 91 * math.sqrt(715) * sh_6_2 * x
+        )
+        sh_7_13 = -3 / 7 * math.sqrt(5) * sh_6_1 * x + (3 / 7) * math.sqrt(5) * sh_6_11 * z + (1 / 7) * math.sqrt(15) * sh_6_12 * y
+        sh_7_14 = (1 / 14) * math.sqrt(210) * (-sh_6_0 * x + sh_6_12 * z)
+        if lmax == 7:
+            return torch.stack(
+                [
+                    sh_1_0,
+                    sh_1_1,
+                    sh_1_2,
+                    sh_2_0,
+                    sh_2_1,
+                    sh_2_2,
+                    sh_2_3,
+                    sh_2_4,
+                    sh_3_0,
+                    sh_3_1,
+                    sh_3_2,
+                    sh_3_3,
+                    sh_3_4,
+                    sh_3_5,
+                    sh_3_6,
+                    sh_4_0,
+                    sh_4_1,
+                    sh_4_2,
+                    sh_4_3,
+                    sh_4_4,
+                    sh_4_5,
+                    sh_4_6,
+                    sh_4_7,
+                    sh_4_8,
+                    sh_5_0,
+                    sh_5_1,
+                    sh_5_2,
+                    sh_5_3,
+                    sh_5_4,
+                    sh_5_5,
+                    sh_5_6,
+                    sh_5_7,
+                    sh_5_8,
+                    sh_5_9,
+                    sh_5_10,
+                    sh_6_0,
+                    sh_6_1,
+                    sh_6_2,
+                    sh_6_3,
+                    sh_6_4,
+                    sh_6_5,
+                    sh_6_6,
+                    sh_6_7,
+                    sh_6_8,
+                    sh_6_9,
+                    sh_6_10,
+                    sh_6_11,
+                    sh_6_12,
+                    sh_7_0,
+                    sh_7_1,
+                    sh_7_2,
+                    sh_7_3,
+                    sh_7_4,
+                    sh_7_5,
+                    sh_7_6,
+                    sh_7_7,
+                    sh_7_8,
+                    sh_7_9,
+                    sh_7_10,
+                    sh_7_11,
+                    sh_7_12,
+                    sh_7_13,
+                    sh_7_14,
+                ],
+                dim=-1,
+            )
+
+        sh_8_0 = (1 / 4) * math.sqrt(17) * (sh_7_0 * z + sh_7_14 * x)
+        sh_8_1 = (1 / 8) * math.sqrt(17) * sh_7_0 * y + (1 / 16) * math.sqrt(238) * sh_7_1 * z + (1 / 16) * math.sqrt(238) * sh_7_13 * x
+        sh_8_2 = (
+            -1 / 240 * math.sqrt(510) * sh_7_0 * z
+            + (1 / 60) * math.sqrt(1785) * sh_7_1 * y
+            + (1 / 240) * math.sqrt(46410) * sh_7_12 * x
+            + (1 / 240) * math.sqrt(510) * sh_7_14 * x
+            + (1 / 240) * math.sqrt(46410) * sh_7_2 * z
+        )
+        sh_8_3 = (
+            (1 / 80)
+            * math.sqrt(2)
+            * (
+                -math.sqrt(85) * sh_7_1 * z
+                + math.sqrt(2210) * sh_7_11 * x
+                + math.sqrt(85) * sh_7_13 * x
+                + math.sqrt(2210) * sh_7_2 * y
+                + math.sqrt(2210) * sh_7_3 * z
+            )
+        )
+        sh_8_4 = (
+            (1 / 40) * math.sqrt(935) * sh_7_10 * x
+            + (1 / 40) * math.sqrt(85) * sh_7_12 * x
+            - 1 / 40 * math.sqrt(85) * sh_7_2 * z
+            + (1 / 10) * math.sqrt(85) * sh_7_3 * y
+            + (1 / 40) * math.sqrt(935) * sh_7_4 * z
+        )
+        sh_8_5 = (
+            (1 / 48)
+            * math.sqrt(2)
+            * (
+                math.sqrt(102) * sh_7_11 * x
+                - math.sqrt(102) * sh_7_3 * z
+                + math.sqrt(1122) * sh_7_4 * y
+                + math.sqrt(561) * sh_7_5 * z
+                + math.sqrt(561) * sh_7_9 * x
+            )
+        )
+        sh_8_6 = (
+            (1 / 16) * math.sqrt(34) * sh_7_10 * x
+            - 1 / 16 * math.sqrt(34) * sh_7_4 * z
+            + (1 / 4) * math.sqrt(17) * sh_7_5 * y
+            + (1 / 16) * math.sqrt(102) * sh_7_6 * z
+            + (1 / 16) * math.sqrt(102) * sh_7_8 * x
+        )
+        sh_8_7 = (
+            -1 / 80 * math.sqrt(1190) * sh_7_5 * z
+            + (1 / 40) * math.sqrt(1785) * sh_7_6 * y
+            + (1 / 20) * math.sqrt(255) * sh_7_7 * x
+            + (1 / 80) * math.sqrt(1190) * sh_7_9 * x
+        )
+        sh_8_8 = -1 / 60 * math.sqrt(1785) * sh_7_6 * x + (1 / 15) * math.sqrt(255) * sh_7_7 * y - 1 / 60 * math.sqrt(1785) * sh_7_8 * z
+        sh_8_9 = (
+            -1 / 80 * math.sqrt(1190) * sh_7_5 * x
+            + (1 / 20) * math.sqrt(255) * sh_7_7 * z
+            + (1 / 40) * math.sqrt(1785) * sh_7_8 * y
+            - 1 / 80 * math.sqrt(1190) * sh_7_9 * z
+        )
+        sh_8_10 = (
+            -1 / 16 * math.sqrt(34) * sh_7_10 * z
+            - 1 / 16 * math.sqrt(34) * sh_7_4 * x
+            - 1 / 16 * math.sqrt(102) * sh_7_6 * x
+            + (1 / 16) * math.sqrt(102) * sh_7_8 * z
+            + (1 / 4) * math.sqrt(17) * sh_7_9 * y
+        )
+        sh_8_11 = (
+            (1 / 48)
+            * math.sqrt(2)
+            * (
+                math.sqrt(1122) * sh_7_10 * y
+                - math.sqrt(102) * sh_7_11 * z
+                - math.sqrt(102) * sh_7_3 * x
+                - math.sqrt(561) * sh_7_5 * x
+                + math.sqrt(561) * sh_7_9 * z
+            )
+        )
+        sh_8_12 = (
+            (1 / 40) * math.sqrt(935) * sh_7_10 * z
+            + (1 / 10) * math.sqrt(85) * sh_7_11 * y
+            - 1 / 40 * math.sqrt(85) * sh_7_12 * z
+            - 1 / 40 * math.sqrt(85) * sh_7_2 * x
+            - 1 / 40 * math.sqrt(935) * sh_7_4 * x
+        )
+        sh_8_13 = (
+            (1 / 80)
+            * math.sqrt(2)
+            * (
+                -math.sqrt(85) * sh_7_1 * x
+                + math.sqrt(2210) * sh_7_11 * z
+                + math.sqrt(2210) * sh_7_12 * y
+                - math.sqrt(85) * sh_7_13 * z
+                - math.sqrt(2210) * sh_7_3 * x
+            )
+        )
+        sh_8_14 = (
+            -1 / 240 * math.sqrt(510) * sh_7_0 * x
+            + (1 / 240) * math.sqrt(46410) * sh_7_12 * z
+            + (1 / 60) * math.sqrt(1785) * sh_7_13 * y
+            - 1 / 240 * math.sqrt(510) * sh_7_14 * z
+            - 1 / 240 * math.sqrt(46410) * sh_7_2 * x
+        )
+        sh_8_15 = -1 / 16 * math.sqrt(238) * sh_7_1 * x + (1 / 16) * math.sqrt(238) * sh_7_13 * z + (1 / 8) * math.sqrt(17) * sh_7_14 * y
+        sh_8_16 = (1 / 4) * math.sqrt(17) * (-sh_7_0 * x + sh_7_14 * z)
+        if lmax == 8:
+            return torch.stack(
+                [
+                    sh_1_0,
+                    sh_1_1,
+                    sh_1_2,
+                    sh_2_0,
+                    sh_2_1,
+                    sh_2_2,
+                    sh_2_3,
+                    sh_2_4,
+                    sh_3_0,
+                    sh_3_1,
+                    sh_3_2,
+                    sh_3_3,
+                    sh_3_4,
+                    sh_3_5,
+                    sh_3_6,
+                    sh_4_0,
+                    sh_4_1,
+                    sh_4_2,
+                    sh_4_3,
+                    sh_4_4,
+                    sh_4_5,
+                    sh_4_6,
+                    sh_4_7,
+                    sh_4_8,
+                    sh_5_0,
+                    sh_5_1,
+                    sh_5_2,
+                    sh_5_3,
+                    sh_5_4,
+                    sh_5_5,
+                    sh_5_6,
+                    sh_5_7,
+                    sh_5_8,
+                    sh_5_9,
+                    sh_5_10,
+                    sh_6_0,
+                    sh_6_1,
+                    sh_6_2,
+                    sh_6_3,
+                    sh_6_4,
+                    sh_6_5,
+                    sh_6_6,
+                    sh_6_7,
+                    sh_6_8,
+                    sh_6_9,
+                    sh_6_10,
+                    sh_6_11,
+                    sh_6_12,
+                    sh_7_0,
+                    sh_7_1,
+                    sh_7_2,
+                    sh_7_3,
+                    sh_7_4,
+                    sh_7_5,
+                    sh_7_6,
+                    sh_7_7,
+                    sh_7_8,
+                    sh_7_9,
+                    sh_7_10,
+                    sh_7_11,
+                    sh_7_12,
+                    sh_7_13,
+                    sh_7_14,
+                    sh_8_0,
+                    sh_8_1,
+                    sh_8_2,
+                    sh_8_3,
+                    sh_8_4,
+                    sh_8_5,
+                    sh_8_6,
+                    sh_8_7,
+                    sh_8_8,
+                    sh_8_9,
+                    sh_8_10,
+                    sh_8_11,
+                    sh_8_12,
+                    sh_8_13,
+                    sh_8_14,
+                    sh_8_15,
+                    sh_8_16,
+                ],
+                dim=-1,
+            )
+
+
+def lmax_tensor_size(lmax):
+    return ((lmax + 1) ** 2) - 1
+
+
+def get_split_sizes_from_dim(feature_dim):
+    """
+    Find the lmax value and return split sizes for torch.split based on feature dimension.
+
+    Args:
+        feature_dim: The dimension of the feature (shape[1] of the tensor)
+
+    Returns:
+        split_sizes: A list of split sizes for torch.split (sizes of spherical harmonic components)
+    """
+    lmax = 1
+    while lmax_tensor_size(lmax) < feature_dim:
+        lmax += 1
+
+    if lmax_tensor_size(lmax) != feature_dim:
+        raise ValueError(f"Feature dimension {feature_dim} does not correspond to a valid lmax value")
+
+    # Return the sizes of each spherical harmonic component
+    return [2 * l + 1 for l in range(1, lmax + 1)]  # noqa: E741
+
+
+class TensorLayerNorm(nn.Module):
+    def __init__(self, hidden_channels, trainable):
+        super(TensorLayerNorm, self).__init__()
+
+        self.hidden_channels = hidden_channels
+        self.eps = 1e-12
+
+        weight = torch.ones(self.hidden_channels)
+        if trainable:
+            self.register_parameter("weight", nn.Parameter(weight))
+        else:
+            self.register_buffer("weight", weight)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        weight = torch.ones(self.hidden_channels)
+        self.weight.data.copy_(weight)
+
+    def max_min_norm(self, tensor):
+        # Based on VisNet (https://www.nature.com/articles/s41467-023-43720-2)
+        dist = torch.norm(tensor, dim=1, keepdim=True)
+
+        if (dist == 0).all():
+            return torch.zeros_like(tensor)
+
+        dist = dist.clamp(min=self.eps)
+        direct = tensor / dist
+
+        max_val, _ = torch.max(dist, dim=-1)
+        min_val, _ = torch.min(dist, dim=-1)
+        delta = (max_val - min_val).view(-1)
+        delta = torch.where(delta == 0, torch.ones_like(delta), delta)
+        dist = (dist - min_val.view(-1, 1, 1)) / delta.view(-1, 1, 1)
+
+        return F.relu(dist) * direct
+
+    def forward(self, tensor):
+        # vec: (num_atoms, feature_dim, hidden_channels)
+        feature_dim = tensor.shape[1]
+
+        try:
+            split_sizes = get_split_sizes_from_dim(feature_dim)
+        except ValueError as e:
+            raise ValueError(f"VecLayerNorm received unsupported feature dimension {feature_dim}: {str(e)}")
+
+        # Split the vector into parts
+        vec_parts = torch.split(tensor, split_sizes, dim=1)
+
+        # Normalize each part separately
+        normalized_parts = [self.max_min_norm(part) for part in vec_parts]
+
+        # Concatenate the normalized parts
+        normalized_vec = torch.cat(normalized_parts, dim=1)
+
+        # Apply weight
+        return normalized_vec * self.weight.unsqueeze(0).unsqueeze(0)
+
+
+def normalize_string(s: str) -> str:
+    return s.lower().replace("-", "").replace("_", "").replace(" ", "")
+
+
+class Swish(nn.Module):
+    def __init__(self):
+        super(Swish, self).__init__()
+
+    def forward(self, x):
+        return x * torch.sigmoid(x)
+
+
+act_class_mapping = {"ssp": ShiftedSoftplus, "silu": nn.SiLU, "tanh": nn.Tanh, "sigmoid": nn.Sigmoid, "swish": Swish}
+
+
+# https://github.com/sunglasses-ai/classy/blob/3e74cba1fdf1b9f9f2ba1cfcfa6c2017aa59fc04/classy/optim/factories.py#L14
+def get_activations(optional=False, *args, **kwargs):
+    activations = {
+        normalize_string(act.__name__): act
+        for act in vars(torch.nn.modules.activation).values()
+        if isinstance(act, type) and issubclass(act, torch.nn.Module)
+    }
+    activations.update(
+        {
+            "relu": torch.nn.ReLU,
+            "elu": torch.nn.ELU,
+            "sigmoid": torch.nn.Sigmoid,
+            "silu": torch.nn.SiLU,
+            "mish": torch.nn.Mish,
+            "swish": torch.nn.SiLU,
+            "selu": torch.nn.SELU,
+            "scaled_swish": scaled_silu,
+            "softplus": shifted_softplus,
+            "slrelu": SmoothLeakyReLU,
+        }
+    )
+
+    if optional:
+        activations[""] = None
+
+    return activations
+
+
+def get_activations_none(optional=False, *args, **kwargs):
+    activations = {
+        normalize_string(act.__name__): act
+        for act in vars(torch.nn.modules.activation).values()
+        if isinstance(act, type) and issubclass(act, torch.nn.Module)
+    }
+    activations.update(
+        {
+            "relu": torch.nn.ReLU,
+            "elu": torch.nn.ELU,
+            "sigmoid": torch.nn.Sigmoid,
+            "silu": torch.nn.SiLU,
+            "selu": torch.nn.SELU,
+        }
+    )
+
+    if optional:
+        activations[""] = None
+        activations[None] = None
+
+    return activations
+
+
+def dictionary_to_option(options, selected):
+    if selected not in options:
+        raise ValueError(f'Invalid choice "{selected}", choose one from {", ".join(list(options.keys()))} ')
+
+    activation = options[selected]
+    if inspect.isclass(activation):
+        activation = activation()
+    return activation
+
+
+def str2act(input_str, *args, **kwargs):
+    if input_str == "":
+        return None
+
+    act = get_activations(optional=True, *args, **kwargs)
+    out = dictionary_to_option(act, input_str)
+    return out
+
+
+class ExpNormalSmearing(nn.Module):
+    def __init__(self, cutoff=5.0, n_rbf=50, trainable=False):
+        super(ExpNormalSmearing, self).__init__()
+        if isinstance(cutoff, torch.Tensor):
+            cutoff = cutoff.item()
+        self.cutoff = cutoff
+        self.n_rbf = n_rbf
+        self.trainable = trainable
+
+        self.cutoff_fn = CosineCutoff(cutoff)
+        self.alpha = 5.0 / cutoff
+
+        means, betas = self._initial_params()
+        if trainable:
+            self.register_parameter("means", nn.Parameter(means))
+            self.register_parameter("betas", nn.Parameter(betas))
+        else:
+            self.register_buffer("means", means)
+            self.register_buffer("betas", betas)
+
+    def _initial_params(self):
+        start_value = torch.exp(torch.scalar_tensor(-self.cutoff))
+        means = torch.linspace(start_value, 1, self.n_rbf)
+        betas = torch.tensor([(2 / self.n_rbf * (1 - start_value)) ** -2] * self.n_rbf)
+        return means, betas
+
+    def reset_parameters(self):
+        means, betas = self._initial_params()
+        self.means.data.copy_(means)
+        self.betas.data.copy_(betas)
+
+    def forward(self, dist):
+        dist = dist.unsqueeze(-1)
+        return self.cutoff_fn(dist) * torch.exp(-self.betas * (torch.exp(self.alpha * (-dist)) - self.means) ** 2)
+
+
+def str2basis(input_str):
+    if type(input_str) != str:  # noqa: E721
+        return input_str
+
+    if input_str == "BesselBasis":
+        radial_basis = BesselBasis
+    elif input_str == "GaussianRBF":
+        radial_basis = GaussianRBF
+    elif input_str.lower() == "expnorm":
+        radial_basis = ExpNormalSmearing
+    else:
+        raise ValueError("Unknown radial basis: {}".format(input_str))
+
+    return radial_basis
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        hidden_dims: List[int],
+        bias=True,
+        activation=None,
+        last_activation=None,
+        weight_init=xavier_uniform_,
+        bias_init=zeros_initializer,
+        norm="",
+    ):
+        super().__init__()
+
+        # hidden_dims = [hidden, half, hidden]
+
+        dims = hidden_dims
+        n_layers = len(dims)
+
+        DenseMLP = partial(Dense, bias=bias, weight_init=weight_init, bias_init=bias_init)
+
+        self.dense_layers = nn.ModuleList(
+            [DenseMLP(dims[i], dims[i + 1], activation=activation, norm=norm) for i in range(n_layers - 2)]
+            + [DenseMLP(dims[-2], dims[-1], activation=last_activation)]
+        )
+
+        self.layers = nn.Sequential(*self.dense_layers)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for m in self.dense_layers:
+            m.reset_parameters()
+
+    def forward(self, x):
+        return self.layers(x)
+
+
+class NodeInit(MessagePassing):
+    def __init__(
+        self,
+        hidden_channels,
+        num_rbf,
+        cutoff,
+        max_z=100,
+        activation=F.silu,
+        proj_ln="",
+        last_activation=False,
+        weight_init=nn.init.xavier_uniform_,
+        bias_init=nn.init.zeros_,
+        concat=False,
+    ):
+        super(NodeInit, self).__init__(aggr="add")
+
+        if type(hidden_channels) == int:  # noqa: E721
+            hidden_channels = [hidden_channels]
+
+        first_channel = hidden_channels[0]
+        last_channel = hidden_channels[-1]
+
+        DenseInit = partial(Dense, weight_init=weight_init, bias_init=bias_init)  # noqa: F841
+
+        self.concat = concat
+        self.embedding = nn.Embedding(max_z, last_channel)
+        if self.concat:
+            self.embedding_src = nn.Embedding(max_z, first_channel)
+            self.distance_proj = MLP(
+                [num_rbf + 2 * first_channel] + hidden_channels,
+                activation=activation,
+                norm=proj_ln,
+                weight_init=weight_init,
+                bias_init=bias_init,
+                last_activation=activation if last_activation else None,
+            )
+        else:
+            self.distance_proj = MLP(
+                [num_rbf] + [last_channel], activation=None, norm="", weight_init=weight_init, bias_init=bias_init, last_activation=None
+            )
+
+        if not self.concat:
+            self.combine = MLP(
+                [2 * last_channel] + hidden_channels,
+                activation=activation,
+                norm=proj_ln,
+                weight_init=weight_init,
+                bias_init=bias_init,
+                last_activation=activation if last_activation else None,
+            )
+        self.cutoff = CosineCutoff(cutoff)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.embedding.reset_parameters()
+        if self.concat:
+            self.embedding_src.reset_parameters()
+        self.distance_proj.reset_parameters()
+        if not self.concat:
+            self.combine.reset_parameters()
+
+    def forward(self, z, x, edge_index, edge_weight, edge_attr):
+        # remove self loops
+        mask = edge_index[0] != edge_index[1]
+        if not mask.all():
+            edge_index = edge_index[:, mask]
+            edge_weight = edge_weight[mask]
+            edge_attr = edge_attr[mask]
+
+        x_neighbors = self.embedding(z)
+        if not self.concat:
+            C = self.cutoff(edge_weight)
+            W = self.distance_proj(edge_attr) * C.view(-1, 1)
+            x_src = x_neighbors
+        else:
+            x_src = self.embedding_src(z)
+            W = edge_attr
+        # propagate_type: (x: Tensor, s:Tensor, W: Tensor)
+        x_neighbors = self.propagate(edge_index, x=x_neighbors, s=x_src, W=W, size=None)
+
+        if self.concat:
+            x_neighbors = x + x_neighbors
+        else:
+            x_neighbors = self.combine(torch.cat([x, x_neighbors], dim=1))
+        return x_neighbors
+
+    def message(self, s_i, x_j, W):
+        if self.concat:
+            return self.distance_proj(torch.cat([W, x_j, s_i], dim=1))
+        return x_j * W
+
+
+class EdgeInit(MessagePassing):
+    def __init__(
+        self,
+        num_rbf,
+        hidden_channels,
+        activation=F.silu,
+        proj_ln="",
+        last_activation=False,
+        weight_init=nn.init.xavier_uniform_,
+        bias_init=nn.init.zeros_,
+    ):
+        super(EdgeInit, self).__init__(aggr=None)
+        self.activation = activation
+
+        if type(hidden_channels) == int:  # noqa: E721
+            hidden_channels = [hidden_channels]
+        self.edge_up = MLP(
+            [num_rbf] + hidden_channels,
+            activation=activation,
+            norm=proj_ln,
+            weight_init=weight_init,
+            bias_init=bias_init,
+            last_activation=activation if last_activation else None,
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        self.edge_up.reset_parameters()
+
+    def forward(self, edge_index, edge_attr, x):
+        # propagate_type: (x: Tensor, edge_attr: Tensor)
+        out = self.propagate(edge_index, x=x, edge_attr=edge_attr)
+        return out
+
+    def message(self, x_i, x_j, edge_attr):
+        return (x_i + x_j) * self.edge_up(edge_attr)
+
+    def aggregate(self, features, index):
+        # no aggregate
+        return features

models/pos_egnn/posegnn/utils.py

@@ -0,0 +1,69 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor, nn
+from torch_nl import compute_neighborlist
+from torch_nl.geometry import compute_distances
+from torch_nl.neighbor_list import compute_cell_shifts
+
+
+ACT_CLASS_MAPPING = {"silu": nn.SiLU, "tanh": nn.Tanh, "sigmoid": nn.Sigmoid, "gelu": nn.GELU}
+
+class BatchedPeriodicDistance(nn.Module):
+    """
+    Wraps the `torch_nl` package to calculate Periodic Distance using
+    PyTorch operations efficiently. Compute the neighbor list for a given cutoff.
+    Reference: https://github.com/felixmusil/torch_nl
+    """
+
+    def __init__(self, cutoff: float = 5.0) -> None:
+        super().__init__()
+        self.cutoff = cutoff
+        self.self_interactions = False
+
+    def forward(
+        self, pos: Tensor, box: Tensor, batch: Optional[Tensor] = None, precomputed_edge_index=None, precomputed_shifts_idx=None
+    ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
+        # No batch, single sample
+        if batch is None:
+            n_atoms = pos.shape[0]
+            batch = torch.zeros(n_atoms, device=pos.device, dtype=torch.int64)
+
+        is_zero = torch.eq(box, 0)
+        is_not_all_zero = ~is_zero.all(dim=-1).all(dim=-1)
+        pbc = is_not_all_zero.unsqueeze(-1).repeat(1, 3)  # We need to change this when dealing with interfaces
+
+        if (precomputed_edge_index is None) or (precomputed_shifts_idx is None):
+            edge_index, batch_mapping, shifts_idx = compute_neighborlist(self.cutoff, pos, box, pbc, batch, self.self_interactions)
+        else:
+            edge_index = precomputed_edge_index
+            shifts_idx = precomputed_shifts_idx
+            batch_mapping = batch[edge_index[0]]  # NOTE: should be same as edge_index[1]
+
+        cell_shifts = compute_cell_shifts(box, shifts_idx, batch_mapping)
+        edge_weight = compute_distances(pos, edge_index, cell_shifts)
+
+        edge_vec = -(pos[edge_index[1]] - pos[edge_index[0]] + cell_shifts)
+
+        # edge_weight and edge_vec should have grad_fn
+        return edge_index, edge_weight, edge_vec, shifts_idx
+
+
+def get_symmetric_displacement(
+    positions: torch.Tensor,
+    box: Optional[torch.Tensor],
+    num_graphs: int,
+    batch: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    displacement = torch.zeros(
+        (num_graphs, 3, 3),
+        dtype=positions.dtype,
+        device=positions.device,
+    )
+    displacement.requires_grad_(True)
+    symmetric_displacement = 0.5 * (displacement + displacement.transpose(-1, -2))
+    positions = positions + torch.einsum("be,bec->bc", positions, symmetric_displacement[batch])
+    box = box.view(-1, 3, 3)
+    box = box + torch.matmul(box, symmetric_displacement)
+
+    return positions, box, displacement

models/pos_egnn/requirements.txt

@@ -0,0 +1,10 @@
+-f https://data.pyg.org/whl/torch-2.5.1%2Bcu121.html
+-f https://data.pyg.org/whl/torch-2.5.1%2Bcpu.html
+numpy==1.26.4
+ase==3.24.0
+torch==2.5.1
+torch_geometric==2.5.3
+torch_nl==0.3
+torch_scatter
+torch_sparse
+tqdm>=4.66.1

requirements.txt

@@ -27,3 +27,5 @@ torch-optimizer
 tqdm>=4.66.4
 pandas==2.2.3
 mordred
+ase==3.24.0
+torch_nl==0.3