From 2c9785d9cad6c34275b2bf9fda8ff8c0a62e31a0 Mon Sep 17 00:00:00 2001
From: Thiago Reschutzegger <thiagoreschutzegger@ibm.com>
Date: Sun, 11 May 2025 22:09:02 +0000
Subject: [PATCH 1/4] One sctructure during inference to support skin

---
 models/pos_egnn/posegnn/calculator.py |  73 ++++++++++++-------
 models/pos_egnn/posegnn/model.py      | 100 +++++++++++++++-----------
 2 files changed, 106 insertions(+), 67 deletions(-)

diff --git a/models/pos_egnn/posegnn/calculator.py b/models/pos_egnn/posegnn/calculator.py
index c185338..140834d 100644
--- a/models/pos_egnn/posegnn/calculator.py
+++ b/models/pos_egnn/posegnn/calculator.py
@@ -1,64 +1,87 @@
+from typing import Optional
+
 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):
+    def __init__(
+        self,
+        checkpoint: str,
+        device: str,
+        compute_stress: bool = True,
+        compile: bool = True,
+        skin: Optional[float] = None,
+        **kwargs,
+    ):
         Calculator.__init__(self, **kwargs)
-
         checkpoint_dict = torch.load(checkpoint, weights_only=True, map_location=device)
 
-        self.model = PosEGNN(checkpoint_dict["config"])
+        self.model = PosEGNN(checkpoint_dict["config"], skin=skin)
         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
 
+        self.implemented_properties = ["energy", "forces"] + (
+            ["stress"] if compute_stress else []
+        )
+
+        if compile:
+            print("Using torch.compile")
+            self.model = torch.compile(
+                self.model, mode="reduce-overhead", fullgraph=True
+            )
+
     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
+        z, pos, box = self._build_tensors(atoms)
+
+        out = self.model.compute_properties(
+            z, pos, box, compute_stress=self.compute_stress
+        )
+
         self.results = {
             "energy": out["total_energy"].cpu().detach().numpy(),
-            "forces": out["force"].cpu().detach().numpy()
+            "forces": out["force"].cpu().detach().numpy(),
         }
+
         if self.compute_stress:
-            self.results.update({
-                "stress": -out["stress"].squeeze().cpu().detach().numpy()
-                })
+            self.results.update(
+                {"stress": -out["stress"].squeeze().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()
+    def _build_tensors(self, atoms: Atoms):
+        atomic_nums = np.array([atomic_numbers[symbol] for symbol in atoms.symbols])
+
+        z = torch.tensor(atomic_nums, 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)
+
+        return z, pos, box
 
 
 def get_invariant_embeddings(self):
     if self.calc is None:
         raise RuntimeError("No calculator is set.")
     else:
-        data = self.calc._build_data(self)
+        z, pos, box = self.calc._build_tensors(self)
         with torch.no_grad():
-            embeddings = self.calc.model(data)["embedding_0"][..., -1].squeeze(2)
+            embeddings, _, _ = self.calc.model(z, pos, box)["embedding_0"][
+                ..., -1
+            ].squeeze(2)
         return embeddings
 
 
diff --git a/models/pos_egnn/posegnn/model.py b/models/pos_egnn/posegnn/model.py
index 940b09f..50614f6 100644
--- a/models/pos_egnn/posegnn/model.py
+++ b/models/pos_egnn/posegnn/model.py
@@ -1,16 +1,25 @@
-from torch import nn
+from typing import Dict
+
 import torch
+from torch import nn, Tensor
 from .encoder import GotenNet
-from .utils import get_symmetric_displacement, BatchedPeriodicDistance, ACT_CLASS_MAPPING
-from torch_scatter import scatter
+from .utils import (
+    get_symmetric_displacement,
+    PeriodicDistance,
+    ACT_CLASS_MAPPING,
+)
+
 
 class NodeInvariantReadout(nn.Module):
-    def __init__(self, in_channels, num_residues, hidden_channels, out_channels, activation):
+    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)])
+        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](),
@@ -18,64 +27,73 @@ def __init__(self, in_channels, num_residues, hidden_channels, out_channels, act
         )
 
     def forward(self, embedding_0):
-        layer_outputs = embedding_0.squeeze(2)  # [n_nodes, in_channels, num_residues]
+        layer_outputs = embedding_0.squeeze(2)
+
+        outputs = torch.stack(
+            [linear(layer_outputs[:, :, i]) for i, linear in enumerate(self.linears)],
+            dim=0,
+        )
 
-        processed_outputs = []
-        for i, linear in enumerate(self.linears):
-            processed_outputs.append(linear(layer_outputs[:, :, i]))
+        last_output = self.non_linear(layer_outputs[:, :, -1]).unsqueeze(0)
+        processed_outputs = torch.cat([outputs, last_output], dim=0)
+        output = processed_outputs.sum(dim=0).squeeze(-1)
 
-        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):
+    def __init__(self, config: Dict, **kwargs):
         super().__init__()
 
-        self.distance = BatchedPeriodicDistance(config["encoder"]["cutoff"])
+        self.distance = PeriodicDistance(
+            config["encoder"]["cutoff"], skin=kwargs.get("skin", None)
+        )
         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)
+        self.register_buffer("e0_mean", Tensor(config["e0_mean"]))
+        self.register_buffer(
+            "atomic_res_total_mean", Tensor(config["atomic_res_total_mean"])
+        )
+        self.register_buffer(
+            "atomic_res_total_std", Tensor(config["atomic_res_total_std"])
+        )
+
+    def forward(self, z: Tensor, pos: Tensor, box: Tensor):
+        pos_grad = pos.clone().requires_grad_(True)
 
-        data.pos, data.box, data.displacements = get_symmetric_displacement(data.pos, data.box, data.num_graphs, data.batch)
+        pos, box, displacements = get_symmetric_displacement(pos_grad, box)
 
-        data.cutoff_edge_index, data.cutoff_edge_distance, data.cutoff_edge_vec, data.cutoff_shifts_idx = self.distance(
-            data.pos, data.box, data.batch
+        cutoff_edge_index, cutoff_edge_distance, cutoff_edge_vec, cutoff_shifts_idx = (
+            self.distance(pos, box)
         )
 
-        embedding_dict = self.encoder(data.z, data.pos, data.cutoff_edge_index, data.cutoff_edge_distance, data.cutoff_edge_vec)
+        embedding_dict = self.encoder(
+            z, pos, cutoff_edge_index, cutoff_edge_distance, cutoff_edge_vec
+        )
 
-        return embedding_dict
+        return embedding_dict, pos, displacements
 
-    def compute_properties(self, data, compute_stress = True):
+    def compute_properties(
+        self, z: Tensor, pos: Tensor, box: Tensor, compute_stress: float = True
+    ):
         output = {}
-        
-        embedding_dict = self.forward(data)
+
+        embedding_dict, pos, displacements = self.forward(z, pos, box)
         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[z]
+        total_energy = node_e0.sum() + node_e_res.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
+            inputs = [pos, displacements]
         else:
-            inputs = [data.pos]
+            inputs = [pos]
 
         grad_outputs = torch.autograd.grad(
             outputs=[total_energy],
@@ -85,15 +103,13 @@ def compute_properties(self, data, compute_stress = True):
             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))
+            stress = virial / torch.det(box).abs().view(-1, 1, 1)
             output["force"] = -force
             output["stress"] = -stress
         else:
             force = grad_outputs[0]
-            output["force"] = -force        
-        
+            output["force"] = -force
+
         return output

From 5433be3630ed72a9a4f9d4e774095595175483b0 Mon Sep 17 00:00:00 2001
From: Thiago Reschutzegger <thiagoreschutzegger@ibm.com>
Date: Sun, 11 May 2025 22:09:30 +0000
Subject: [PATCH 2/4] Skin implementation

---
 models/pos_egnn/posegnn/utils.py | 102 +++++++++++++++++++++----------
 1 file changed, 70 insertions(+), 32 deletions(-)

diff --git a/models/pos_egnn/posegnn/utils.py b/models/pos_egnn/posegnn/utils.py
index 7153562..72c7bac 100644
--- a/models/pos_egnn/posegnn/utils.py
+++ b/models/pos_egnn/posegnn/utils.py
@@ -1,60 +1,96 @@
 from typing import Optional, Tuple
 
 import torch
-from torch import Tensor, nn
+from torch import 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}
+ACT_CLASS_MAPPING = {
+    "silu": nn.SiLU,
+    "tanh": nn.Tanh,
+    "sigmoid": nn.Sigmoid,
+    "gelu": nn.GELU,
+}
 
-class BatchedPeriodicDistance(nn.Module):
+
+class PeriodicDistance(nn.Module):
     """
     Wraps the `torch_nl` package to calculate Periodic Distance using
-    PyTorch operations efficiently. Compute the neighbor list for a given cutoff.
+    PyTorch operations efficiently. Compute the neighbor list for a given cutoff
+    with an optional 'skin' buffer to cache results until atoms move
+    beyond a threshold.
     Reference: https://github.com/felixmusil/torch_nl
     """
 
-    def __init__(self, cutoff: float = 5.0) -> None:
+    def __init__(self, cutoff: float = 6.0, skin: Optional[float] = None) -> None:
         super().__init__()
         self.cutoff = cutoff
+        self.skin = skin
         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)
+        if skin is not None:
+            self.register_buffer("_prev_pos", None)
+            self.register_buffer("_prev_box", None)
+            self._cached_topo = None
+
+    def forward(self, pos: torch.Tensor, box: torch.Tensor):
+        # This method supports one structure only
+        batch = torch.zeros(pos.size(0), dtype=torch.long, device=pos.device)
+
+        rebuild = True
+        if self.skin is not None:
+            rebuild = self._cached_topo is None or self._moved_beyond_skin(pos, box)
 
+        if rebuild:
+            is_zero = torch.eq(box, 0.0)
+            pbc_mask = ~is_zero.all(dim=-1).all(dim=-1)
+            pbc = pbc_mask.unsqueeze(-1).repeat(1, 3)
+
+            edge_index, batch_map, shifts_idx = compute_neighborlist(
+                self.cutoff, pos, box, pbc, batch, self.self_interactions
+            )
+            self._cached_topo = (edge_index, shifts_idx, batch_map)
+            self._prev_pos = pos.clone()
+            self._prev_box = box.clone()
+
+        edge_index, shifts_idx, batch_map = self._cached_topo
+
+        cell_shifts = compute_cell_shifts(box, shifts_idx, batch_map)
+        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 _moved_beyond_skin(self, pos: torch.Tensor, box: torch.Tensor):
+        prev_box = self._prev_box[0]
+
+        if not torch.allclose(box, prev_box, rtol=1e-5, atol=1e-8):  # Box deformation
+            return True
+
+        if torch.allclose(box, torch.zeros_like(box)):
+            delta = pos - self._prev_pos
+            max_disp = torch.linalg.norm(delta, dim=1).max()
+            return max_disp.item() > self.skin
+
+        inv_box, _ = torch.linalg.inv_ex(prev_box, check_errors=True)
+        delta = pos.unsqueeze(0) - self._prev_pos.unsqueeze(0)
+        frac = torch.einsum("ij,bkj->bki", inv_box, delta)
+        frac = frac - torch.round(frac)
+        delta = torch.einsum("bij,bkj->bki", self._prev_box, frac)
+        max_disp = torch.linalg.norm(delta.squeeze(0), dim=1).max()
+
+        return max_disp.item() > self.skin
+
 
 def get_symmetric_displacement(
-    positions: torch.Tensor,
-    box: Optional[torch.Tensor],
-    num_graphs: int,
-    batch: torch.Tensor,
+    positions: torch.Tensor, box: Optional[torch.Tensor]
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    # This method supports one structure only
+    batch = torch.zeros(positions.size(0), dtype=torch.long, device=positions.device)
+    num_graphs = 1
+
     displacement = torch.zeros(
         (num_graphs, 3, 3),
         dtype=positions.dtype,
@@ -62,7 +98,9 @@ def get_symmetric_displacement(
     )
     displacement.requires_grad_(True)
     symmetric_displacement = 0.5 * (displacement + displacement.transpose(-1, -2))
-    positions = positions + torch.einsum("be,bec->bc", positions, symmetric_displacement[batch])
+    positions = positions + torch.einsum(
+        "be,bec->bc", positions, symmetric_displacement[batch]
+    )
     box = box.view(-1, 3, 3)
     box = box + torch.matmul(box, symmetric_displacement)
 

From 04bef12bcf2b392fba02b5dd4d25574d9bc5e52a Mon Sep 17 00:00:00 2001
From: Thiago Reschutzegger <thiagoreschutzegger@ibm.com>
Date: Sun, 11 May 2025 22:09:45 +0000
Subject: [PATCH 3/4] Ruff format

---
 models/pos_egnn/load.py             |  23 ++--
 models/pos_egnn/posegnn/__init__.py |   2 +-
 models/pos_egnn/posegnn/encoder.py  |  95 +++++++++++----
 models/pos_egnn/posegnn/ops.py      | 181 ++++++++++++++++++++++------
 4 files changed, 234 insertions(+), 67 deletions(-)

diff --git a/models/pos_egnn/load.py b/models/pos_egnn/load.py
index 7cf2127..b72333e 100644
--- a/models/pos_egnn/load.py
+++ b/models/pos_egnn/load.py
@@ -1,28 +1,26 @@
 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()
+        numbers=[atom.GetAtomicNum() for atom in mol.GetAtoms()],
+        positions=mol.GetConformer().GetPositions(),
     )
     return ase_atoms
 
-class POSEGNN():
+
+class POSEGNN:
     def __init__(self, use_gpu=True):
         device = "cuda" if use_gpu and torch.cuda.is_available() else "cpu"
         self.device = device
@@ -32,14 +30,16 @@ 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)
+        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]
+            batch = smiles_list[i : i + batch_size]
             atoms_batch = []
 
             for smiles in batch:
@@ -51,7 +51,9 @@ def encode(self, smiles_list, return_tensor=False, batch_size=32):
                     print(f"Skipping {smiles}: {e}")
 
             if atoms_batch:
-                embeddings = [a.get_invariant_embeddings().mean(dim=0).cpu() for a in 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)
 
@@ -60,4 +62,3 @@ def encode(self, smiles_list, return_tensor=False, batch_size=32):
 
         all_embeddings = torch.cat(results, dim=0)
         return all_embeddings if return_tensor else pd.DataFrame(all_embeddings.numpy())
-
diff --git a/models/pos_egnn/posegnn/__init__.py b/models/pos_egnn/posegnn/__init__.py
index 9fb1220..2ac5233 100644
--- a/models/pos_egnn/posegnn/__init__.py
+++ b/models/pos_egnn/posegnn/__init__.py
@@ -1,3 +1,3 @@
 from . import calculator, encoder, model, ops, utils
 
-__all__ = ["calculator", "encoder", "model", "ops", "utils"]
\ No newline at end of file
+__all__ = ["calculator", "encoder", "model", "ops", "utils"]
diff --git a/models/pos_egnn/posegnn/encoder.py b/models/pos_egnn/posegnn/encoder.py
index 7c70bbc..0b5abb4 100644
--- a/models/pos_egnn/posegnn/encoder.py
+++ b/models/pos_egnn/posegnn/encoder.py
@@ -45,7 +45,9 @@ def split_degree(tensor, lmax, dim=-1):  # default to last dim
         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
+        slc[dim] = slice(
+            cumsum, cumsum + count
+        )  # Replace desired dim with actual slice
         tensors.append(tensor[tuple(slc)])
         cumsum += count
     return tensors
@@ -129,16 +131,28 @@ def __init__(
             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
+                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
             )
-            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)]
+                    [
+                        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)
+                self.veck_w = InitDense(
+                    n_atom_basis, self.edge_vec_dim, activation=None, bias=False
+                )
 
             if self.update_info["lin_w"] > 0:
                 modules = []
@@ -148,7 +162,9 @@ def __init__(
                     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
+                    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)
@@ -254,13 +270,17 @@ def forward(
             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)
+                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 = 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
@@ -292,24 +312,32 @@ def message(
         Compute message passing.
         """
 
-        r_ij_attn = r_ij_attn.reshape(-1, self.num_heads, self.n_atom_basis // self.num_heads)
+        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)
+            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)
+        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)))
+        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
@@ -379,10 +407,14 @@ def aggregate(
     ) -> 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)
+        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]:
+    def update(
+        self, inputs: Tuple[torch.Tensor, torch.Tensor]
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         return inputs
 
 
@@ -422,7 +454,9 @@ 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)
+        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)
@@ -440,7 +474,7 @@ class GotenNet(nn.Module):
     def __init__(
         self,
         hidden_channels: int = 128,
-        num_layers: int = 8, 
+        num_layers: int = 8,
         radial_basis: Union[Callable, str] = "BesselBasis",
         n_rbf: int = 20,
         cutoff: float = 5.0,
@@ -496,7 +530,11 @@ def __init__(
             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=""
+            n_rbf,
+            [self.hidden_dim // 2, self.hidden_dim],
+            weight_init=weight_init,
+            bias_init=bias_init,
+            proj_ln="",
         )
 
         radial_basis = str2basis(radial_basis)
@@ -535,7 +573,13 @@ def __init__(
 
         self.eqff = nn.ModuleList(
             [
-                EQFF(n_atom_basis=self.n_atom_basis, activation=activation, epsilon=epsilon, weight_init=weight_init, bias_init=bias_init)
+                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)
             ]
         )
@@ -571,7 +615,9 @@ def forward(self, z, pos, cutoff_edge_index, cutoff_edge_distance, cutoff_edge_v
 
         edge_attr = self.radial_basis(cutoff_edge_distance)
 
-        q = self.neighbor_embedding(z, q, cutoff_edge_index, cutoff_edge_distance, edge_attr)
+        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
@@ -581,7 +627,12 @@ def forward(self, z, pos, cutoff_edge_index, cutoff_edge_distance, cutoff_edge_v
         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")
+        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]]
@@ -615,7 +666,9 @@ def forward(self, z, pos, cutoff_edge_index, cutoff_edge_distance, cutoff_edge_v
         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]
+        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
diff --git a/models/pos_egnn/posegnn/ops.py b/models/pos_egnn/posegnn/ops.py
index c3f5160..03cbc30 100644
--- a/models/pos_egnn/posegnn/ops.py
+++ b/models/pos_egnn/posegnn/ops.py
@@ -23,7 +23,7 @@
 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
+from torch_scatter import scatter
 
 zeros_initializer = partial(constant_, val=0.0)
 
@@ -34,7 +34,9 @@ def centralize(
     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_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
@@ -64,9 +66,21 @@ def parse_update_info(edge_updates):
     else:
         update_parts = []
 
-    allowed_parts = ["gated", "gatedt", "norej", "mlp", "mlpa", "act", "linw", "linwa", "drej"]
+    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}")
+        raise ValueError(
+            f"Invalid edge update parts. Allowed parts are {allowed_parts}"
+        )
 
     if "gated" in update_parts:
         update_info["gated"] = "gated"
@@ -228,7 +242,9 @@ class SchnetMLP(nn.Module):
             any activation function.
     """
 
-    def __init__(self, n_in, n_out, n_hidden=None, n_layers=2, activation=shifted_softplus):
+    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:
@@ -245,7 +261,10 @@ def __init__(self, n_in, n_out, n_hidden=None, n_layers=2, activation=shifted_so
             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)]
+        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
@@ -275,7 +294,9 @@ def gaussian_rbf(inputs: torch.Tensor, offsets: torch.Tensor, widths: torch.Tens
 class GaussianRBF(nn.Module):
     r"""Gaussian radial basis functions."""
 
-    def __init__(self, n_rbf: int, cutoff: float, start: float = 0.0, trainable: bool = False):
+    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`.
@@ -289,7 +310,9 @@ def __init__(self, n_rbf: int, cutoff: float, start: float = 0.0, trainable: boo
 
         # 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))
+        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)
@@ -535,7 +558,9 @@ def __init__(self, l=2):  # noqa: E741
         self.l = l
 
     def forward(self, edge_vec):
-        edge_sh = self._calculate_components(self.l, edge_vec[..., 0], edge_vec[..., 1], edge_vec[..., 2])
+        edge_sh = self._calculate_components(
+            self.l, edge_vec[..., 0], edge_vec[..., 1], edge_vec[..., 2]
+        )
         return edge_sh
 
     @property
@@ -543,7 +568,9 @@ 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:
+    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:
@@ -560,7 +587,9 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
         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)
+            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)
@@ -594,7 +623,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             )
 
         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_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
@@ -608,7 +641,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             + (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_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
@@ -622,7 +659,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             + (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_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(
@@ -656,7 +697,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             )
 
         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_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
@@ -677,7 +722,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             + (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_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
@@ -698,7 +747,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             + (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_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(
@@ -743,7 +796,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             )
 
         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_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
@@ -771,7 +828,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             + (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_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
@@ -799,7 +860,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             + (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_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(
@@ -857,7 +922,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             )
 
         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_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
@@ -892,7 +961,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             + (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_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
@@ -927,7 +1000,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             - 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_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(
@@ -1000,7 +1077,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             )
 
         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_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
@@ -1050,7 +1131,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             + (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_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
@@ -1100,7 +1185,11 @@ def _calculate_components(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.
             - 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_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(
@@ -1209,7 +1298,9 @@ def get_split_sizes_from_dim(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")
+        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
@@ -1259,7 +1350,9 @@ def forward(self, tensor):
         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)}")
+            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)
@@ -1286,7 +1379,13 @@ 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}
+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
@@ -1342,7 +1441,9 @@ def get_activations_none(optional=False, *args, **kwargs):
 
 def dictionary_to_option(options, selected):
     if selected not in options:
-        raise ValueError(f'Invalid choice "{selected}", choose one from {", ".join(list(options.keys()))} ')
+        raise ValueError(
+            f'Invalid choice "{selected}", choose one from {", ".join(list(options.keys()))} '
+        )
 
     activation = options[selected]
     if inspect.isclass(activation):
@@ -1392,7 +1493,9 @@ def reset_parameters(self):
 
     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)
+        return self.cutoff_fn(dist) * torch.exp(
+            -self.betas * (torch.exp(self.alpha * (-dist)) - self.means) ** 2
+        )
 
 
 def str2basis(input_str):
@@ -1429,10 +1532,15 @@ def __init__(
         dims = hidden_dims
         n_layers = len(dims)
 
-        DenseMLP = partial(Dense, bias=bias, weight_init=weight_init, bias_init=bias_init)
+        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[i], dims[i + 1], activation=activation, norm=norm)
+                for i in range(n_layers - 2)
+            ]
             + [DenseMLP(dims[-2], dims[-1], activation=last_activation)]
         )
 
@@ -1486,7 +1594,12 @@ def __init__(
             )
         else:
             self.distance_proj = MLP(
-                [num_rbf] + [last_channel], activation=None, norm="", weight_init=weight_init, bias_init=bias_init, last_activation=None
+                [num_rbf] + [last_channel],
+                activation=None,
+                norm="",
+                weight_init=weight_init,
+                bias_init=bias_init,
+                last_activation=None,
             )
 
         if not self.concat:

From 28bca7e0b62c2b04249ec78b902de8a92430e813 Mon Sep 17 00:00:00 2001
From: Thiago Reschutzegger <thiagoreschutzegger@ibm.com>
Date: Sun, 11 May 2025 22:38:43 +0000
Subject: [PATCH 4/4] Undo substitution

---
 models/pos_egnn/posegnn/model.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/models/pos_egnn/posegnn/model.py b/models/pos_egnn/posegnn/model.py
index 50614f6..b2311ed 100644
--- a/models/pos_egnn/posegnn/model.py
+++ b/models/pos_egnn/posegnn/model.py
@@ -51,12 +51,12 @@ def __init__(self, config: Dict, **kwargs):
         self.encoder = GotenNet(**config["encoder"])
         self.readout = NodeInvariantReadout(**config["decoder"])
 
-        self.register_buffer("e0_mean", Tensor(config["e0_mean"]))
+        self.register_buffer("e0_mean", torch.tensor(config["e0_mean"]))
         self.register_buffer(
-            "atomic_res_total_mean", Tensor(config["atomic_res_total_mean"])
+            "atomic_res_total_mean", torch.tensor(config["atomic_res_total_mean"])
         )
         self.register_buffer(
-            "atomic_res_total_std", Tensor(config["atomic_res_total_std"])
+            "atomic_res_total_std", torch.tensor(config["atomic_res_total_std"])
         )
 
     def forward(self, z: Tensor, pos: Tensor, box: Tensor):