add test (#6)

Author: MatteoMazzonelli

.pre-commit-config.yaml

@@ -20,11 +20,6 @@ repos:
     rev: v3.2.0
     hooks:
     -   id: add-trailing-comma
--   repo: https://github.com/PyCQA/flake8
-    rev: 7.2.0
-    hooks:
-    -   id: flake8
-        args: ["--ignore=E501,W503,E203"]
 -   repo: https://github.com/ibm/detect-secrets
     rev: 0.13.1+ibm.62.dss
     hooks:

tests/create_reference.py

@@ -0,0 +1,36 @@
+# save_reference_output.py
+import torch
+import numpy as np
+
+torch.manual_seed(0)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+num_nodes = 1
+x_dim = 9
+pe_dim = 20
+edge_attr_dim = 2
+
+# Dummy all-zero input
+x = torch.zeros((num_nodes, x_dim)).to(device)
+pe = torch.zeros((num_nodes, pe_dim)).to(device)
+edge_index = torch.tensor([[0], [0]]).to(device)
+edge_attr = torch.zeros((1, edge_attr_dim)).to(device)
+batch = torch.zeros(num_nodes, dtype=torch.long).to(device)
+
+
+models = {
+    "v0_1_2": "../examples/models/GridFM_v0_1_2.pth",
+    "v0_2_3": "../examples/models/GridFM_v0_2_3.pth",
+}
+
+for version, path in models.items():
+    print(f"Loading model {version}...")
+    model = torch.load(path, weights_only=False, map_location=device).to(device)
+    model.eval()
+
+    with torch.no_grad():
+        output = model(x, pe, edge_index, edge_attr, batch)
+
+    out_path = f"./data/reference_output_{version}.npy"
+    np.save(out_path, output.cpu().numpy())
+    print(f"Saved output for {version} to {out_path}")

tests/data/reference_output_v0_1_2.npy

@@ -0,0 +1,55 @@
+import torch
+import numpy as np
+import pytest
+
+# Device setup
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Input shape config
+num_nodes = 1
+x_dim = 9
+pe_dim = 20
+edge_attr_dim = 2
+
+# List of models and reference files to check
+models_to_test = [
+    (
+        "v0_1_2",
+        "examples/models/GridFM_v0_1_2.pth",
+        "tests/data/reference_output_v0_1_2.npy",
+    ),
+    (
+        "v0_2_3",
+        "examples/models/GridFM_v0_2_3.pth",
+        "tests/data/reference_output_v0_2_3.npy",
+    ),
+]
+
+
+@pytest.mark.parametrize("version, model_path, ref_output_path", models_to_test)
+def test_model_matches_reference(version, model_path, ref_output_path):
+    torch.manual_seed(0)
+
+    # Prepare zero input
+    x = torch.zeros((num_nodes, x_dim), device=device)
+    pe = torch.zeros((num_nodes, pe_dim), device=device)
+    edge_index = torch.tensor([[0], [0]], device=device)
+    edge_attr = torch.zeros((1, edge_attr_dim), device=device)
+    batch = torch.zeros(num_nodes, dtype=torch.long, device=device)
+
+    # Load model
+    model = torch.load(model_path, weights_only=False, map_location=device).to(device)
+    model.eval()
+
+    # Get current output
+    with torch.no_grad():
+        output = model(x, pe, edge_index, edge_attr, batch).cpu().numpy()
+
+    # Load saved reference
+    reference = np.load(ref_output_path)
+
+    # Exact match assertion
+    assert np.allclose(output, reference, rtol=1e-5, atol=1e-6), (
+        f"Model output for {version} does not match reference within tolerance.\n"
+        f"Max absolute difference: {np.max(np.abs(output - reference))}"
+    )