Dynamic Node Labels (#386)

* WIP

* Update

* fix existing tests

* update examples

* wip

* update docs

* update tgbn label semantics

* add some tests

* more tests

* more tests

* permissions

* coverage

* spelling

* Address comments

Author: Jacob-Chmura

README.md

@@ -135,19 +135,19 @@ class RecurrentGCN(torch.nn.Module):
 
 # Initialize our model and optimizer
 encoder = RecurrentGCN(node_dim=static_node_x.shape[1], embed_dim=128)
-decoder = NodePredictor(in_dim=128, out_dim=train_dg.node_x_dim)
+decoder = NodePredictor(in_dim=128, out_dim=train_dg.node_y_dim)
 opt = torch.optim.Adam(set(encoder.parameters()) | set(decoder.parameters()), lr=0.001)
 
 # Training loop
 h_0 = None
 for batch in train_loader:
     opt.zero_grad()
-    y_true = batch.node_x
+    y_true = batch.node_y
     if y_true is None:
         continue
 
     z, h_0 = encoder(batch, static_node_x, h_0)
-    z_node = z[batch.node_x_nids]
+    z_node = z[batch.node_y_nids]
     y_pred = decoder(z_node)
 
     loss = F.cross_entropy(y_pred, y_true)

docs/quickstart.md

@@ -40,19 +40,19 @@ class RecurrentGCN(torch.nn.Module):
 
 # Initialize our model and optimizer
 encoder = RecurrentGCN(node_dim=static_node_x.shape[1], embed_dim=128)
-decoder = NodePredictor(in_dim=128, out_dim=train_dg.node_x_dim)
+decoder = NodePredictor(in_dim=128, out_dim=train_dg.node_y_dim)
 opt = torch.optim.Adam(set(encoder.parameters()) | set(decoder.parameters()), lr=0.001)
 
 # Training loop
 h_0 = None
 for batch in train_loader:
     opt.zero_grad()
-    y_true = batch.node_x
+    y_true = batch.node_y
     if y_true is None:
         continue
 
     z, h_0 = encoder(batch, static_node_x, h_0)
-    z_node = z[batch.node_x_nids]
+    z_node = z[batch.node_y_nids]
     y_pred = decoder(z_node)
 
     loss = F.cross_entropy(y_pred, y_true)

docs/tutorials/dgraph_tutorial.md

@@ -54,19 +54,27 @@ class DGData:
         edge_mask (Tensor): Mask of edge events within `time`.
         edge_index (Tensor): Edge connections [num_edge_events, 2].
         edge_x (Tensor | None): Optional edge features [num_edge_events, D_edge].
-        node_mask (Tensor | None): Mask of node events within `time`.
-        node_x_nids (Tensor | None): Node IDs corresponding to node events [num_node_events].
-        node_x (Tensor | None): Node features over time [num_node_events, D_node_dynamic].
+        node_x_mask (Tensor | None): Mask of dynamic node features within `time`.
+        node_x_nids (Tensor | None): Node IDs corresponding to dynamic node features [num_node_events].
+        node_x (Tensor | None): Dynamic Node features over time [num_node_events, D_node_dynamic].
+        node_y_mask (Tensor | None): Mask of node labels within `time`.
+        node_y_nids (Tensor | None): Node IDs corresponding to node labels [num_node_labels].
+        node_y (Tensor | None): Node labels over time [num_node_labels, D_node_dynamic].
         static_node_x (Tensor | None): Node features invariant over time [num_nodes, D_node_static].
+        edge_type (Tensor | None) : Type of relation of each edge event in edge_index [num_edge_events].
+        node_type (Tensor | None) : Type of each node [num_nodes].
 
     Raises:
         InvalidNodeIDError: If an edge or node ID match `PADDED_NODE_ID`.
+        InvalidNodeIDError: If node labels exists with node IDs outside the graph's node ID range.
         ValueError: If any data attributes have non-well defined tensor shapes.
         EmptyGraphError: If attempting to initialize an empty graph.
 
     Notes:
         - Timestamps must be non-negative and sorted; DGData will sort automatically if necessary.
         - Cloning creates a deep copy of tensors to prevent in-place modifications.
+        - Edge type is only applicable for Heterogeneous & Knowledge graph.
+        - Node type is only applicable for Knowledge graph.
     """
 ```
 
@@ -114,10 +122,14 @@ Please consult our documentation for full description of our API. The table belo
 | `edge_src_col`            | Column name in edge file for src nodes                         | `str`                 | Yes                                               | Cannot have ids matching `tgm.constants.PADDED_NODE_ID`  |
 | `edge_dst_col`            | Column name in edge file for dst nodes                         | `str`                 | Yes                                               | Cannot have ids matching  `tgm.constants.PADDED_NODE_ID` |
 | `edge_time_col`           | Column name in edge file for edge times                        | `str`                 | Yes                                               | Time must be non-negative                                |
-| `node_file_path`          | Path to CSV file containing dynamic node data                  | `str \| pathlib.Path` | No                                                | `node_df` is using `from_pandas`                         |
-| `node_x_nids_col`         | Column name in node file for node event node ids               | `str`                 | No, unless `node_file_path` is specified          | Cannot have ids matching  `tgm.constants.PADDED_NODE_ID` |
-| `node_x_time_col`         | Column name in node file for node event node times             | `str`                 | No, unless `node_file_path` is specified          | Time must be non-negative                                |
+| `node_x_file_path`        | Path to CSV file containing dynamic node data                  | `str \| pathlib.Path` | No                                                | `node_x_df` is using `from_pandas`                       |
+| `node_x_nids_col`         | Column name in node file for node event node ids               | `str`                 | No, unless `node_x_file_path` is specified        | Cannot have ids matching  `tgm.constants.PADDED_NODE_ID` |
+| `node_x_time_col`         | Column name in node file for node event node times             | `str`                 | No, unless `node_x_file_path` is specified        | Time must be non-negative                                |
 | `node_x_col`              | Column name in node file for dynamic node features             | `str`                 | No                                                |                                                          |
+| `node_y_file_path`        | Path to CSV file containing dynamic node labels                | `str \| pathlib.Path` | No                                                | `node_y_df` is using `from_pandas`                       |
+| `node_y_nids_col`         | Column name in node file for node label node ids               | `str`                 | No, unless `node_y_file_path` is specified        | Cannot have ids matching  `tgm.constants.PADDED_NODE_ID` |
+| `node_y_time_col`         | Column name in node file for node label node times             | `str`                 | No, unless `node_y_file_path` is specified        | Time must be non-negative                                |
+| `node_y_col`              | Column name in node file for dynamic node labels               | `str`                 | No                                                |                                                          |
 | `static_node_x_file_path` | Path to CSV file containing static node features               | `str \| pathlib.Path` | No                                                | `static_node_x_df` if using `from_pandas`                |
 | `static_node_x_col`       | Column name in static node feats file for static node features | `str`                 | No, unless `static_node_x_file_path` is specified |                                                          |
 | `time_delta`              | Time granularity of the graph data                             | `TimeDeltaDG \| str`  | Yes                                               | Default to *event_ordered* granularity `'r'`             |
@@ -131,16 +143,18 @@ A few key things to know:
   - We expect an `edge_file_path` which is a csv file with `edge_src_col`, `edge_dst_col`, `edge_time_col` as a minimum.
   - Your edge csv file may also contain `edge_x_col` which are the edge features on your data
 - dynamic node data (optional)
-  - If included, we expect a `node_file_path` which is a csv file with `node_x_nids_col`, `node_x_time_col` as a minimum. These are your dynamic node events.
+  - If included, we expect a `node_x_file_path` which is a csv file with `node_x_nids_col`, `node_x_time_col` as a minimum. These are your dynamic node events.
+  - If included, we expect a `node_y_file_path` which is a csv file with `node_y_nids_col`, `node_y_time_col` as a minimum. These are your dynamic node labels.
   - Your dynamic node data csv file may also include `node_x_col`, which are the dynamic node features in your data.
+  - Your dynamic node labels csv file may also include `node_y_col`, which are the dynamic node labels in your data.
 - static node data (optional)
   - If included, we expect a `static_node_x_file_path` which is a csv file with `static_node_x_col`, the static node features for your dataset.
 
 Internally, we perform various checks on the tensors shapes, node ranges, and timestamps values. If your data is well structured, everything should work. If you get an error message that is not intuitive, please let us know.
 
 #### From Pandas
 
-The API largely the same as above, except that we expected `edge_df`, `node_df`, and `static_node_x_df` dataframes for the edge, dynamic node, and static node data respectively, instead of csv files.
+The API largely the same as above, except that we expected `edge_df`, `node_x_df`, and `static_node_x_df` dataframes for the edge, dynamic node, and static node data respectively, instead of csv files.
 
 ```python
 import pandas as pd
@@ -171,7 +185,7 @@ dg = DGData.from_pandas(
     edge_dst_col='dst',
     edge_time_col='t',
     edge_x_col='edge_feat',
-    node_df=dynamic_node_df,
+    node_x_df=dynamic_node_df,
     node_x_nids_col='node',
     node_x_time_col='t',
     node_x_col='dynamic_node_feat',
@@ -404,6 +418,9 @@ class DGBatch:
         node_x (Tensor | None, optional): Dynamic node features for nodes in the batch. Tensor of shape `(T x V x d_node_dynamic)`.
         node_x_time (Tensor | None, optional): Timestamps corresponding to dynamic node features.
         node_x_nids (Tensor | None, optional): Node IDs corresponding to dynamic node features.
+        node_y (Tensor | None, optional): Dynamic node labels for nodes in the batch. Tensor of shape `(T x V x d_node_labels)`.
+        node_y_time (Tensor | None, optional): Timestamps corresponding to dynamic node labels.
+        node_y_nids (Tensor | None, optional): Node IDs corresponding to dynamic node labels.
     """
 ```
 

docs/tutorials/hook_tutorial.md

@@ -48,7 +48,7 @@ class MyNegativeHook(StatelessHook):
     requires = set()
 
     def __call__(self, dg: DGraph, batch: DGBatch) -> DGBatch:
-        batch.my_neg = torch.randint(10, 100, (len(batch.dst),))
+        batch.my_neg = torch.randint(10, 100, (len(batch.edge_dst),))
         batch.my_neg_time = batch.edge_time.clone()
         return batch
 ```

examples/nodeproppred/dygformer.py

@@ -176,7 +176,7 @@ def train(
     for batch in tqdm(loader):
         opt.zero_grad()
 
-        y_true = batch.node_x
+        y_true = batch.node_y
         if len(batch.edge_src) > 0:
             z = encoder(batch, static_node_x)  # [num_nodes, embed_dim]
 
@@ -210,12 +210,12 @@ def eval(
     static_node_x = loader.dgraph.static_node_x
 
     for batch in tqdm(loader):
-        y_true = batch.node_x
+        y_true = batch.node_y
 
         if batch.edge_src.shape[0] > 0:
             z = encoder(batch, static_node_x)
             if y_true is not None:
-                z_node = z[batch.node_x_nids]
+                z_node = z[batch.node_y_nids]
                 y_pred = decoder(z_node)
                 input_dict = {
                     'y_true': y_true,
@@ -263,7 +263,7 @@ def eval(
 val_loader = DGDataLoader(val_dg, batch_size=args.bsize, hook_manager=hm)
 test_loader = DGDataLoader(test_dg, batch_size=args.bsize, hook_manager=hm)
 
-num_classes = train_dg.node_x_dim
+num_classes = train_dg.node_y_dim
 
 encoder = DyGFormer_NodePrediction(
     num_nodes=full_data.num_nodes,

examples/nodeproppred/gclstm.py

@@ -81,12 +81,12 @@ def train(
 
     for batch in tqdm(loader):
         opt.zero_grad()
-        y_true = batch.node_x
+        y_true = batch.node_y
         if y_true is None:
             continue
 
         z, h_0, c_0 = encoder(batch, static_node_x, h_0, c_0)
-        z_node = z[batch.node_x_nids]
+        z_node = z[batch.node_y_nids]
         y_pred = decoder(z_node)
 
         loss = F.cross_entropy(y_pred, y_true)
@@ -116,12 +116,12 @@ def eval(
     static_node_x = loader.dgraph.static_node_x
 
     for batch in tqdm(loader):
-        y_true = batch.node_x
+        y_true = batch.node_y
         if y_true is None:
             continue
 
         z, h_0, c_0 = encoder(batch, static_node_x, h_0, c_0)
-        z_node = z[batch.node_x_nids]
+        z_node = z[batch.node_y_nids]
         y_pred = decoder(z_node)
 
         input_dict = {
@@ -152,7 +152,7 @@ def eval(
 val_loader = DGDataLoader(val_dg, batch_unit=args.snapshot_time_gran)
 test_loader = DGDataLoader(test_dg, batch_unit=args.snapshot_time_gran)
 
-num_classes = train_dg.node_x_dim
+num_classes = train_dg.node_y_dim
 
 encoder = RecurrentGCN(
     node_dim=train_dg.static_node_x_dim, embed_dim=args.embed_dim

examples/nodeproppred/gcn.py

@@ -100,12 +100,12 @@ def train(
 
     for batch in tqdm(loader):
         opt.zero_grad()
-        y_true = batch.node_x
+        y_true = batch.node_y
         if y_true is None:
             continue
 
         z = encoder(batch, static_node_x)
-        z_node = z[batch.node_x_nids]
+        z_node = z[batch.node_y_nids]
         y_pred = decoder(z_node)
 
         loss = F.cross_entropy(y_pred, y_true)
@@ -131,12 +131,12 @@ def eval(
     static_node_x = loader.dgraph.static_node_x
 
     for batch in tqdm(loader):
-        y_true = batch.node_x
+        y_true = batch.node_y
         if y_true is None:
             continue
 
         z = encoder(batch, static_node_x)
-        z_node = z[batch.node_x_nids]
+        z_node = z[batch.node_y_nids]
         y_pred = decoder(z_node)
 
         input_dict = {
@@ -167,7 +167,7 @@ def eval(
 val_loader = DGDataLoader(val_dg, batch_unit=args.snapshot_time_gran)
 test_loader = DGDataLoader(test_dg, batch_unit=args.snapshot_time_gran)
 
-num_classes = train_dg.node_x_dim
+num_classes = train_dg.node_y_dim
 
 encoder = GCNEncoder(
     in_channels=train_dg.static_node_x_dim,

examples/nodeproppred/persistant_forecast.py

@@ -52,12 +52,12 @@ def eval(
     perf_list = []
 
     for batch in tqdm(loader):
-        y_true = batch.node_x
+        y_true = batch.node_y
         if y_true is None:
             continue
 
         y_pred = torch.zeros_like(y_true)
-        for i, node_id in enumerate(batch.node_x_nids.tolist()):
+        for i, node_id in enumerate(batch.node_y_nids.tolist()):
             y_pred[i] = model(node_id)
             model.update(node_id, y_true[i])
 
@@ -83,7 +83,7 @@ def eval(
 val_loader = DGDataLoader(val_dg, batch_unit=args.snapshot_time_gran)
 test_loader = DGDataLoader(test_dg, batch_unit=args.snapshot_time_gran)
 
-num_classes = train_dg.node_x_dim
+num_classes = train_dg.node_y_dim
 model = PersistantForecaster(num_classes=num_classes)
 
 eval(train_loader, model, evaluator)

examples/nodeproppred/tgat.py

@@ -144,7 +144,7 @@ def train(
 
     for batch in tqdm(loader):
         opt.zero_grad()
-        y_labels = batch.node_x
+        y_labels = batch.node_y
         if y_labels is not None:
             z = encoder(batch, static_node_x)
             y_pred = decoder(z)
@@ -172,7 +172,7 @@ def eval(
     static_node_x = loader.dgraph.static_node_x
 
     for batch in tqdm(loader):
-        y_labels = batch.node_x
+        y_labels = batch.node_y
         if y_labels is not None:
             z = encoder(batch, static_node_x)
             y_pred = decoder(z)
@@ -198,20 +198,20 @@ def eval(
 val_dg = DGraph(val_data, device=args.device)
 test_dg = DGraph(test_data, device=args.device)
 
-num_classes = train_dg.node_x_dim
+num_classes = train_dg.node_y_dim
 
 if args.sampling == 'uniform':
     nbr_hook = NeighborSamplerHook(
         num_nbrs=args.n_nbrs,
-        seed_nodes_keys=['node_x_nids'],
-        seed_times_keys=['node_x_time'],
+        seed_nodes_keys=['node_y_nids'],
+        seed_times_keys=['node_y_time'],
     )
 elif args.sampling == 'recency':
     nbr_hook = RecencyNeighborHook(
         num_nbrs=args.n_nbrs,
         num_nodes=full_data.num_nodes,  # Assuming node ids at test set > train/val set
-        seed_nodes_keys=['node_x_nids'],
-        seed_times_keys=['node_x_time'],
+        seed_nodes_keys=['node_y_nids'],
+        seed_times_keys=['node_y_time'],
     )
 else:
     raise ValueError(f'Unknown sampling type: {args.sampling}')
@@ -225,8 +225,6 @@ def eval(
 val_loader = DGDataLoader(val_dg, args.bsize, hook_manager=hm)
 test_loader = DGDataLoader(test_dg, args.bsize, hook_manager=hm)
 
-num_classes = train_dg.node_x_dim
-
 encoder = TGAT(
     node_dim=train_dg.static_node_x_dim,
     edge_dim=train_dg.edge_x_dim,

examples/nodeproppred/tgcn.py

@@ -80,12 +80,12 @@ def train(
 
     for batch in tqdm(loader):
         opt.zero_grad()
-        y_true = batch.node_x
+        y_true = batch.node_y
         if y_true is None:
             continue
 
         z, h_0 = encoder(batch, static_node_x)
-        z_node = z[batch.node_x_nids]
+        z_node = z[batch.node_y_nids]
         y_pred = decoder(z_node)
 
         loss = F.cross_entropy(y_pred, y_true)
@@ -114,12 +114,12 @@ def eval(
     static_node_x = loader.dgraph.static_node_x
 
     for batch in tqdm(loader):
-        y_true = batch.node_x
+        y_true = batch.node_y
         if y_true is None:
             continue
 
         z, h_0 = encoder(batch, static_node_x, h_0)
-        z_node = z[batch.node_x_nids]
+        z_node = z[batch.node_y_nids]
         y_pred = decoder(z_node)
 
         input_dict = {
@@ -150,7 +150,7 @@ def eval(
 val_loader = DGDataLoader(val_dg, batch_unit=args.snapshot_time_gran)
 test_loader = DGDataLoader(test_dg, batch_unit=args.snapshot_time_gran)
 
-num_classes = train_dg.node_x_dim
+num_classes = train_dg.node_y_dim
 
 encoder = RecurrentGCN(
     node_dim=train_dg.static_node_x_dim, embed_dim=args.embed_dim