Ruff for Linting and Formatting (#104)

* adding pre-commit and linting / formatting via ruff

* linting + remove transpose permutation in forward

* readme tutorial formatted. Can exclude if preferred

* ruff format and lint

* only test deterministic if shared_weights != true

* revert Tensor -> tensor change (Tensor is a type, tensor the constructor)

* add ruff and pre-commit to the CI requirements .txt (exact version to promote caching)

* skips and looser thresholds

* save failed tensors

* remove tensor saving

* make E741 (lowercase L) lint a package level preference

* add pre-commit github action

* lint and format

---------

Co-authored-by: Austin Glover <austin_glover@berkeley.com>

Author: asglover

.github/workflows/pre-commit.yaml

@@ -0,0 +1,13 @@
+name: Pre-Commit Checks
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: pre-commit/action@v3.0.1

.github/workflows/requirements_cuda_ci.txt

@@ -1,4 +1,6 @@
 numpy==2.2.5
 torch==2.7.0 --index-url https://download.pytorch.org/whl/cu128
 pytest==8.3.5
-ninja==1.11.1.4
+ninja==1.11.1.4
+ruff==0.11.11
+pre-commit==4.2.0

.pre-commit-config.yaml

@@ -0,0 +1,9 @@
+repos:
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  # Ruff version.
+  rev: v0.11.11
+  hooks:
+    # Run the linter.
+    - id: ruff-check
+    # Run the formatter.
+    - id: ruff-format

examples/readme_tutorial.py

@@ -1,26 +1,29 @@
-# Examples from the README 
+# ruff: noqa: E402
+# Examples from the README
 import logging
 from openequivariance.benchmark.logging_utils import getLogger
+
 logger = getLogger()
 logger.setLevel(logging.ERROR)
 
-# UVU Tensor Product 
+# UVU Tensor Product
 # ===============================
 import torch
 import e3nn.o3 as o3
 
-gen = torch.Generator(device='cuda')
+gen = torch.Generator(device="cuda")
 
 batch_size = 1000
-X_ir, Y_ir, Z_ir = o3.Irreps("1x2e"), o3.Irreps("1x3e"), o3.Irreps("1x2e") 
-X = torch.rand(batch_size, X_ir.dim, device='cuda', generator=gen)
-Y = torch.rand(batch_size, Y_ir.dim, device='cuda', generator=gen)
+X_ir, Y_ir, Z_ir = o3.Irreps("1x2e"), o3.Irreps("1x3e"), o3.Irreps("1x2e")
+X = torch.rand(batch_size, X_ir.dim, device="cuda", generator=gen)
+Y = torch.rand(batch_size, Y_ir.dim, device="cuda", generator=gen)
 
-instructions=[(0, 0, 0, "uvu", True)]
+instructions = [(0, 0, 0, "uvu", True)]
 
-tp_e3nn = o3.TensorProduct(X_ir, Y_ir, Z_ir, instructions,
-        shared_weights=False, internal_weights=False).to('cuda')
-W = torch.rand(batch_size, tp_e3nn.weight_numel, device='cuda', generator=gen)
+tp_e3nn = o3.TensorProduct(
+    X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False
+).to("cuda")
+W = torch.rand(batch_size, tp_e3nn.weight_numel, device="cuda", generator=gen)
 
 Z = tp_e3nn(X, Y, W)
 print(torch.norm(Z))
@@ -29,10 +32,12 @@
 # ===============================
 import openequivariance as oeq
 
-problem = oeq.TPProblem(X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False)
+problem = oeq.TPProblem(
+    X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False
+)
 tp_fast = oeq.TensorProduct(problem, torch_op=True)
 
-Z = tp_fast(X, Y, W) # Reuse X, Y, W from earlier
+Z = tp_fast(X, Y, W)  # Reuse X, Y, W from earlier
 print(torch.norm(Z))
 # ===============================
 
@@ -44,26 +49,35 @@
 
 # Receiver, sender indices for message passing GNN
 edge_index = EdgeIndex(
-                [[0, 1, 1, 2],  # Receiver 
-                 [1, 0, 2, 1]], # Sender 
-                device='cuda',
-                dtype=torch.long)
-
-X = torch.rand(node_ct, X_ir.dim, device='cuda', generator=gen)
-Y = torch.rand(nonzero_ct, Y_ir.dim, device='cuda', generator=gen)
-W = torch.rand(nonzero_ct, problem.weight_numel, device='cuda', generator=gen)
-
-tp_conv = oeq.TensorProductConv(problem, torch_op=True, deterministic=False) # Reuse problem from earlier
-Z = tp_conv.forward(X, Y, W, edge_index[0], edge_index[1]) # Z has shape [node_ct, z_ir.dim]
+    [
+        [0, 1, 1, 2],  # Receiver
+        [1, 0, 2, 1],  # Sender
+    ],
+    device="cuda",
+    dtype=torch.long,
+)
+
+X = torch.rand(node_ct, X_ir.dim, device="cuda", generator=gen)
+Y = torch.rand(nonzero_ct, Y_ir.dim, device="cuda", generator=gen)
+W = torch.rand(nonzero_ct, problem.weight_numel, device="cuda", generator=gen)
+
+tp_conv = oeq.TensorProductConv(
+    problem, torch_op=True, deterministic=False
+)  # Reuse problem from earlier
+Z = tp_conv.forward(
+    X, Y, W, edge_index[0], edge_index[1]
+)  # Z has shape [node_ct, z_ir.dim]
 print(torch.norm(Z))
 # ===============================
 
 # ===============================
-_, sender_perm = edge_index.sort_by("col")            # Sort by sender index 
-edge_index, receiver_perm = edge_index.sort_by("row") # Sort by receiver index
+_, sender_perm = edge_index.sort_by("col")  # Sort by sender index
+edge_index, receiver_perm = edge_index.sort_by("row")  # Sort by receiver index
 
 # Now we can use the faster deterministic algorithm
-tp_conv = oeq.TensorProductConv(problem, torch_op=True, deterministic=True) 
-Z = tp_conv.forward(X, Y[receiver_perm], W[receiver_perm], edge_index[0], edge_index[1], sender_perm) 
+tp_conv = oeq.TensorProductConv(problem, torch_op=True, deterministic=True)
+Z = tp_conv.forward(
+    X, Y[receiver_perm], W[receiver_perm], edge_index[0], edge_index[1], sender_perm
+)
 print(torch.norm(Z))
-# ===============================
+# ===============================

io/cif_to_graph.py

@@ -1,42 +1,44 @@
 import pickle
 import numpy as np
 from sklearn.neighbors import radius_neighbors_graph
-from scipy.io import mmwrite
+
 
 def cif_to_molecular_graph(cif_file, cp, radii):
-    with open(f'../data/cif_files/{cif_file}', 'r') as f:
+    with open(f"../data/cif_files/{cif_file}", "r") as f:
         print("Started reading file...")
         lines = f.readlines()
         print("Finished reading file!")
 
         coords = []
         for line in lines:
-            if line.startswith('ATOM'):
+            if line.startswith("ATOM"):
                 parts = line.split()
-                coords.append([float(parts[cp[0]]), float(parts[cp[1]]), float(parts[cp[2]])])
+                coords.append(
+                    [float(parts[cp[0]]), float(parts[cp[1]]), float(parts[cp[2]])]
+                )
 
         coords = np.array(coords)
 
         for radius in radii:
             print(f"Starting radius neighbors calculation, r={radius}")
-            A = radius_neighbors_graph(coords, radius, mode='connectivity',
-                            include_self=False) 
-            print(f"Finished radius neighbors calculation, found {A.nnz} nonzeros.") 
-    
+            A = radius_neighbors_graph(
+                coords, radius, mode="connectivity", include_self=False
+            )
+            print(f"Finished radius neighbors calculation, found {A.nnz} nonzeros.")
+
             # mmwrite(f'../data/molecular_structures/{cif_file.split(".")[0]}.mtx', A)
 
             coo_mat = A.tocoo()
-            result = {
-                'row': coo_mat.row,
-                'col': coo_mat.col,
-                'coords': coords
-            }
+            result = {"row": coo_mat.row, "col": coo_mat.col, "coords": coords}
 
-            with open(f'../data/molecular_structures/{cif_file.split(".")[0]}_radius{radius}.pickle', 'wb') as handle:
+            with open(
+                f"../data/molecular_structures/{cif_file.split('.')[0]}_radius{radius}.pickle",
+                "wb",
+            ) as handle:
                 pickle.dump(result, handle, protocol=pickle.HIGHEST_PROTOCOL)
 
 
-if __name__=='__main__':
-    #cif_to_molecular_graph('hiv_capsid.cif', (10, 11, 12), radii=[2.0, 2.5, 3.0, 3.5])  
-    cif_to_molecular_graph('covid_spike.cif', (10, 11, 12), radii=[2.0, 2.5, 3.0, 3.5])  
-    cif_to_molecular_graph('1drf.cif', (10, 11, 12), radii=[6.0])  
+if __name__ == "__main__":
+    # cif_to_molecular_graph('hiv_capsid.cif', (10, 11, 12), radii=[2.0, 2.5, 3.0, 3.5])
+    cif_to_molecular_graph("covid_spike.cif", (10, 11, 12), radii=[2.0, 2.5, 3.0, 3.5])
+    cif_to_molecular_graph("1drf.cif", (10, 11, 12), radii=[6.0])

io/load_nequip_configs.py

@@ -1,27 +1,29 @@
-'''
+"""
 This script parse the repository of
 Nequip input files at
 https://github.com/mir-group/nequip-input-files.
 We extract the node / edge hidden features representations.
-'''
+"""
+
+import os
+import yaml
 
-import os, yaml
 
 def process_nequip_configs():
     nequip_files = []
-    for root, dirs, files in os.walk('../data/nequip-input-files'):
+    for root, dirs, files in os.walk("../data/nequip-input-files"):
         for file in files:
-            if file.endswith('.yaml'):
+            if file.endswith(".yaml"):
                 nequip_files.append(os.path.join(root, file))
- 
+
     irrep_pairs = []
     configs = []
     for file in nequip_files:
-        with open(file, 'r') as f:
+        with open(file, "r") as f:
             data = yaml.unsafe_load(f)
             filename = os.path.splitext(os.path.basename(file))[0]
-            feature_irreps_hidden = data['feature_irreps_hidden']
-            irreps_edge_sh = data['irreps_edge_sh']
+            feature_irreps_hidden = data["feature_irreps_hidden"]
+            irreps_edge_sh = data["irreps_edge_sh"]
             if (feature_irreps_hidden, irreps_edge_sh) not in irrep_pairs:
                 irrep_pairs.append((feature_irreps_hidden, irreps_edge_sh))
                 configs.append((feature_irreps_hidden, irreps_edge_sh, filename))
@@ -30,5 +32,5 @@ def process_nequip_configs():
         print(config)
 
 
-if __name__ == '__main__':
-    process_nequip_configs()
+if __name__ == "__main__":
+    process_nequip_configs()

openequivariance/__init__.py

@@ -1,18 +1,32 @@
+# ruff: noqa: F401
 import openequivariance.extlib
 from pathlib import Path
 from importlib.metadata import version
 
 from openequivariance.implementations.e3nn_lite import TPProblem, Irreps
-from openequivariance.implementations.TensorProduct import TensorProduct 
-from openequivariance.implementations.convolution.TensorProductConv import TensorProductConv 
+from openequivariance.implementations.TensorProduct import TensorProduct
+from openequivariance.implementations.convolution.TensorProductConv import (
+    TensorProductConv,
+)
 from openequivariance.implementations.utils import torch_to_oeq_dtype
 
+__all__ = [
+    "TPProblem",
+    "Irreps",
+    "TensorProduct",
+    "TensorProductConv",
+    "torch_to_oeq_dtype",
+]
+
 __version__ = version("openequivariance")
 
+
 def _check_package_editable():
     import json
     from importlib.metadata import Distribution
+
     direct_url = Distribution.from_name("openequivariance").read_text("direct_url.json")
     return json.loads(direct_url).get("dir_info", {}).get("editable", False)
 
-_editable_install_output_path = Path(__file__).parent.parent / "outputs"
+
+_editable_install_output_path = Path(__file__).parent.parent / "outputs"