[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

examples/relbench_example.py

@@ -26,8 +26,7 @@
 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)
@@ -41,10 +40,8 @@
 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.
@@ -70,9 +67,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()
@@ -141,21 +138,18 @@ def test() -> Tuple[float, float, float]:
     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

@@ -69,8 +69,7 @@ def from_relbench(db: Any) -> 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 feature_cols:
@@ -84,17 +83,16 @@ def from_relbench(db: Any) -> 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 = 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():
@@ -104,8 +102,7 @@ def from_relbench(db: Any) -> 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