Address akihironitta review: annotations, idioms, style, drop string forward ref

Author: AJamal27891

examples/relbench_example.py

@@ -14,7 +14,9 @@
     ``python relbench_example.py``
     ``python relbench_example.py --epochs 50 --hidden_channels 128``
 """
+
 import argparse
+from typing import Tuple
 
 import torch
 import torch.nn.functional as F
@@ -24,7 +26,8 @@
 from torch_geometric.utils import from_relbench
 
 parser = argparse.ArgumentParser(
-    description='Train a heterogeneous GNN on a RelBench dataset.')
+    description='Train a heterogeneous GNN on a RelBench dataset.'
+)
 parser.add_argument('--hidden_channels', type=int, default=64)
 parser.add_argument('--lr', type=float, default=0.005)
 parser.add_argument('--epochs', type=int, default=30)
@@ -38,14 +41,16 @@
 dataset = get_dataset('rel-f1', download=True)
 db = dataset.get_db()
 data = from_relbench(db)
-print(f'Graph: {len(data.node_types)} node types, '
-      f'{len(data.edge_types)} edge types')
+print(
+    f'Graph: {len(data.node_types)} node types, '
+    f'{len(data.edge_types)} edge types'
+)
 
 # 2. Prepare a node regression target.
 # `from_relbench` preserves the original DataFrame column order from RelBench.
 # In rel-f1, the 'standings' table has 'points' as its first numeric column:
 target_type = 'standings'
-y = data[target_type].x[:, 0].clone()  # points column (index 0 in rel-f1)
+y = data[target_type].x[:, 0]  # points column (index 0 in rel-f1)
 data[target_type].x = data[target_type].x[:, 1:]  # remove from input features
 
 # 3. Clean up features — fill NaN and standardize per column:
@@ -65,9 +70,9 @@
 train_mask = torch.zeros(num_nodes, dtype=torch.bool)
 val_mask = torch.zeros(num_nodes, dtype=torch.bool)
 test_mask = torch.zeros(num_nodes, dtype=torch.bool)
-train_mask[perm[:int(0.6 * num_nodes)]] = True
-val_mask[perm[int(0.6 * num_nodes):int(0.8 * num_nodes)]] = True
-test_mask[perm[int(0.8 * num_nodes):]] = True
+train_mask[perm[: int(0.6 * num_nodes)]] = True
+val_mask[perm[int(0.6 * num_nodes) : int(0.8 * num_nodes)]] = True
+test_mask[perm[int(0.8 * num_nodes) :]] = True
 
 # Normalize target using training set statistics only (prevents data leakage):
 y_mean = y[train_mask].mean()
@@ -92,7 +97,11 @@ def __init__(self, hidden_channels: int) -> None:
         self.conv2 = SAGEConv((-1, -1), hidden_channels)
         self.lin = Linear(-1, 1)
 
-    def forward(self, x, edge_index):
+    def forward(
+        self,
+        x: torch.Tensor,
+        edge_index: torch.Tensor,
+    ) -> torch.Tensor:
         x = self.conv1(x, edge_index).relu()
         x = self.conv2(x, edge_index).relu()
         return self.lin(x)
@@ -108,40 +117,45 @@ def forward(self, x, edge_index):
 optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
 
 
-def train() -> float:
+def train() -> torch.Tensor:
     model.train()
     optimizer.zero_grad()
     pred = model(data.x_dict, data.edge_index_dict)[target_type].squeeze(-1)
     loss = F.mse_loss(pred[train_mask], y_norm[train_mask])
     loss.backward()
     optimizer.step()
-    return float(loss)
+    return loss
 
 
 @torch.no_grad()
-def test():
+def test() -> Tuple[float, float, float]:
     model.eval()
     pred = model(data.x_dict, data.edge_index_dict)[target_type].squeeze(-1)
-    pred_orig = pred * y_std + y_mean  # denormalize for interpretable MAE
+    # denormalize for interpretable MAE
+    pred *= y_std
+    pred += y_mean
 
-    train_mae = float((pred_orig[train_mask] - y[train_mask]).abs().mean())
-    val_mae = float((pred_orig[val_mask] - y[val_mask]).abs().mean())
-    test_mae = float((pred_orig[test_mask] - y[test_mask]).abs().mean())
+    train_mae = float((pred[train_mask] - y[train_mask]).abs().mean())
+    val_mae = float((pred[val_mask] - y[val_mask]).abs().mean())
+    test_mae = float((pred[test_mask] - y[test_mask]).abs().mean())
     return train_mae, val_mae, test_mae
 
 
-print(
-    f'\nTraining {args.epochs} epochs on "{target_type}" point prediction...')
+print(f'\nTraining {args.epochs} epochs on "{target_type}" point prediction...')
 print(f'Target stats (train): mean={y_mean:.2f}, std={y_std:.2f}\n')
 
 for epoch in range(1, args.epochs + 1):
     loss = train()
     if epoch % 5 == 0 or epoch == 1:
         train_mae, val_mae, test_mae = test()
-        print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, '
-              f'Train MAE: {train_mae:.2f}, Val MAE: {val_mae:.2f}, '
-              f'Test MAE: {test_mae:.2f} points')
+        print(
+            f'Epoch: {epoch:03d}, Loss: {loss:.4f}, '
+            f'Train MAE: {train_mae:.2f}, Val MAE: {val_mae:.2f}, '
+            f'Test MAE: {test_mae:.2f} points'
+        )
 
 train_mae, val_mae, test_mae = test()
-print(f'\nFinal — Train MAE: {train_mae:.2f}, Val MAE: {val_mae:.2f}, '
-      f'Test MAE: {test_mae:.2f} points')
+print(
+    f'\nFinal — Train MAE: {train_mae:.2f}, Val MAE: {val_mae:.2f}, '
+    f'Test MAE: {test_mae:.2f} points'
+)

torch_geometric/utils/relbench.py

@@ -3,12 +3,11 @@
 import numpy as np
 import torch
 
-import torch_geometric
 from torch_geometric.data import HeteroData
 from torch_geometric.utils import sort_edge_index
 
 
-def from_relbench(db: Any) -> 'torch_geometric.data.HeteroData':
+def from_relbench(db: Any) -> HeteroData:
     r"""Converts a :class:`relbench.base.Database` object into a
     :class:`~torch_geometric.data.HeteroData` object.
 
@@ -29,7 +28,7 @@ def from_relbench(db: Any) -> 'torch_geometric.data.HeteroData':
         type and each foreign key relationship maps to a pair of directed
         edge types.
 
-    Example:
+    Examples:
         >>> from relbench.base import Database, Table
         >>> import pandas as pd
         >>> users = Table(
@@ -70,10 +69,11 @@ def from_relbench(db: Any) -> 'torch_geometric.data.HeteroData':
 
         # Convert numeric feature columns into a node feature tensor:
         feature_cols = [
-            col for col in df.columns
+            col
+            for col in df.columns
             if col not in exclude_cols and df[col].dtype.kind in ('i', 'f')
         ]
-        if len(feature_cols) > 0:
+        if feature_cols:
             x_np = df[feature_cols].to_numpy(
                 dtype=np.float32,
                 na_value=np.nan,
@@ -84,16 +84,17 @@ def from_relbench(db: Any) -> 'torch_geometric.data.HeteroData':
         if table.time_col is not None:
             time_ser = df[table.time_col]
             if time_ser.dtype in [
-                    np.dtype("datetime64[s]"),
-                    np.dtype("datetime64[ns]"),
+                np.dtype('datetime64[s]'),
+                np.dtype('datetime64[ns]'),
             ]:
-                unix_time = time_ser.astype("int64").values
-                if time_ser.dtype == np.dtype("datetime64[ns]"):
+                unix_time = time_ser.astype('int64').values
+                if time_ser.dtype == np.dtype('datetime64[ns]'):
                     unix_time = unix_time // 10**9
                 data[table_name].time = torch.from_numpy(unix_time)
             else:
                 data[table_name].time = torch.from_numpy(
-                    time_ser.values.astype(np.float64), )
+                    time_ser.values.astype(np.float64)
+                )
 
         # Create edges from foreign key relationships:
         for fkey_col, pkey_table_name in table.fkey_col_to_pkey_table.items():
@@ -103,17 +104,18 @@ def from_relbench(db: Any) -> 'torch_geometric.data.HeteroData':
             mask = ~pkey_index.isna()
             fkey_idx = torch.arange(len(pkey_index))
             pkey_idx = torch.from_numpy(
-                pkey_index[mask].to_numpy(dtype=np.int64), )
+                pkey_index[mask].to_numpy(dtype=np.int64)
+            )
             fkey_idx = fkey_idx[torch.from_numpy(mask.to_numpy(dtype=bool))]
 
             # Forward edge: fkey table -> pkey table
             edge_index = torch.stack([fkey_idx, pkey_idx], dim=0)
-            edge_type = (table_name, f"f2p_{fkey_col}", pkey_table_name)
+            edge_type = (table_name, f'f2p_{fkey_col}', pkey_table_name)
             data[edge_type].edge_index = sort_edge_index(edge_index)
 
             # Reverse edge: pkey table -> fkey table
             edge_index = torch.stack([pkey_idx, fkey_idx], dim=0)
-            edge_type = (pkey_table_name, f"rev_f2p_{fkey_col}", table_name)
+            edge_type = (pkey_table_name, f'rev_f2p_{fkey_col}', table_name)
             data[edge_type].edge_index = sort_edge_index(edge_index)
 
     data.validate()