Repository: pyg-team/pytorch_geometric
Pull request #8454 (2023-11-27 15:55:57)

Merge `TemporalData` with `Data`/`HeteroData` [1/n]
This PR attempts to merge TemporalData into Data.
Currently, only edge events are supported. For `(src, dst)` event definition `edge_index` is reused as its definition is aligned with our needs. For `msg`, `edge_attr` is reused. An additional attribute is `time`, which represents an event's occurrence time.
New methods include the following:
- `is_sorted_by_time()`: returns `True` if data is sorted according to `time`, `False` otherwise.
- `sort_by_time()`: sorts event data according to `time`.
- `snapshot(start_time, end_time)`: returns event data that occurred in range `<start_time, end_time>`.
- `up_to(time)`: returns event data that occurred before or at a given `time`.
- `concat(data)`: returns new data that concatenates the object on which the method was called and given `data`. Can be used to add new events.

Related issue: #3230 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8454?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: `9 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`64fc4c1`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/64fc4c1a58fd661401747db881c366185bcf4912?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.74% compared to head [(`c17b155`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8454?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.74%.

> :exclamation: Current head c17b155 differs from pull request most recent head 08b5b41. Consider uploading reports for the commit 08b5b41 to get more accurate results

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8454?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/data/storage.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8454?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2RhdGEvc3RvcmFnZS5weQ==) | 84.74% | [9 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8454?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8454   +/-   ##
=======================================
  Coverage   88.74%   88.74%           
=======================================
  Files         479      479           
  Lines       29842    29928   +86     
=======================================
+ Hits        26482    26559   +77     
- Misses       3360     3369    +9     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8454?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e6136c611..d5585e2ce 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added temporal-related capabilities to `Data` ([#8454](https://github.com/pyg-team/pytorch_geometric/pull/8454))
 - Added support for returning multi graphs in `to_networkx` ([#8575](https://github.com/pyg-team/pytorch_geometric/pull/8575))
 - Added support for XPU device in `profileit` decorator ([#8532](https://github.com/pyg-team/pytorch_geometric/pull/8532))
 - Added `KNNIndex` exclusion logic ([#8573](https://github.com/pyg-team/pytorch_geometric/pull/8573))
diff --git a/torch_geometric/data/data.py b/torch_geometric/data/data.py
index 8d5dcc1b4..3a6ce9bc1 100644
--- a/torch_geometric/data/data.py
+++ b/torch_geometric/data/data.py
@@ -98,6 +98,15 @@ class BaseData:
         """
         raise NotImplementedError
 
+    def concat(self, data: 'BaseData') -> 'BaseData':
+        r"""Concatenates :obj:`self` with another :obj:`data` object.
+        All values needs to have matching shapes at non-concat dimensions.
+        """
+        out = copy.copy(self)
+        for store, other_store in zip(out.stores, data.stores):
+            store.concat(other_store)
+        return out
+
     def __cat_dim__(self, key: str, value: Any, *args, **kwargs) -> Any:
         r"""Returns the dimension for which the value :obj:`value` of the
         attribute :obj:`key` will get concatenated when creating mini-batches
@@ -256,6 +265,39 @@ class BaseData:
             store.coalesce()
         return out
 
+    def is_sorted_by_time(self) -> bool:
+        r"""Returns :obj:`True` if :obj:`time` is sorted."""
+        return all([store.is_sorted_by_time() for store in self.stores])
+
+    def sort_by_time(self) -> 'BaseData':
+        r"""Sorts data associated with :obj:`time` according to :obj:`time`."""
+        out = copy.copy(self)
+        for store in out.stores:
+            store.sort_by_time()
+        return out
+
+    def snapshot(
+        self,
+        start_time: Union[float, int],
+        end_time: Union[float, int],
+    ) -> 'BaseData':
+        r"""Returns a snapshot of :obj:`data` to only hold events that occurred
+        in period :obj:`[start_time, end_time]`.
+        """
+        out = copy.copy(self)
+        for store in out.stores:
+            store.snapshot(start_time, end_time)
+        return out
+
+    def up_to(self, end_time: Union[float, int]) -> 'BaseData':
+        r"""Returns a snapshot of :obj:`data` to only hold events that occurred
+        up to :obj:`end_time` (inclusive of :obj:`edge_time`).
+        """
+        out = copy.copy(self)
+        for store in out.stores:
+            store.up_to(end_time)
+        return out
+
     def has_isolated_nodes(self) -> bool:
         r"""Returns :obj:`True` if the graph contains isolated nodes."""
         return any([store.has_isolated_nodes() for store in self.edge_stores])
@@ -457,6 +499,8 @@ class Data(BaseData, FeatureStore, GraphStore):
             labels with arbitrary shape. (default: :obj:`None`)
         pos (torch.Tensor, optional): Node position matrix with shape
             :obj:`[num_nodes, num_dimensions]`. (default: :obj:`None`)
+        time (torch.Tensor, optional): The timestamps for each event with shape
+            :obj:`[num_edges]` or :obj:`[num_nodes]`. (default: :obj:`None`)
         **kwargs (optional): Additional attributes.
     """
     def __init__(
@@ -466,6 +510,7 @@ class Data(BaseData, FeatureStore, GraphStore):
         edge_attr: OptTensor = None,
         y: OptTensor = None,
         pos: OptTensor = None,
+        time: OptTensor = None,
         **kwargs,
     ):
         # `Data` doesn't support group_name, so we need to adjust `TensorAttr`
@@ -488,6 +533,8 @@ class Data(BaseData, FeatureStore, GraphStore):
             self.y = y
         if pos is not None:
             self.pos = pos
+        if time is not None:
+            self.time = time
 
         for key, value in kwargs.items():
             setattr(self, key, value)
@@ -944,6 +991,14 @@ class Data(BaseData, FeatureStore, GraphStore):
     def batch(self, batch: Optional[Tensor]):
         self._store.batch = batch
 
+    @property
+    def time(self) -> Optional[Tensor]:
+        return self['time'] if 'time' in self._store else None
+
+    @time.setter
+    def time(self, time: Optional[Tensor]):
+        self._store.time = time
+
     # Deprecated functions ####################################################
 
     @property
diff --git a/torch_geometric/data/storage.py b/torch_geometric/data/storage.py
index edf90816d..6f6b0a42c 100644
--- a/torch_geometric/data/storage.py
+++ b/torch_geometric/data/storage.py
@@ -33,6 +33,7 @@ from torch_geometric.utils import (
     contains_isolated_nodes,
     is_torch_sparse_tensor,
     is_undirected,
+    select,
     sort_edge_index,
 )
 
@@ -286,6 +287,95 @@ class BaseStorage(MutableMapping):
         """
         return self.apply_(lambda x: x.record_stream(stream), *args)
 
+    # Time Handling ###########################################################
+
+    def _cat_dims(self, keys: List[int]) -> Dict[str, int]:
+        return {
+            key: self._parent().__cat_dim__(key, self[key], self)
+            for key in keys
+        }
+
+    def _select(
+        self,
+        keys: List[str],
+        index_or_mask: Tensor,
+    ) -> 'BaseStorage':
+
+        for key, dim in self._cat_dims(keys).items():
+            self[key] = select(self[key], index_or_mask, dim)
+
+        return self
+
+    def concat(self, other: 'GlobalStorage') -> 'GlobalStorage':
+        if not (set(self.keys()) == set(other.keys())):
+            raise AttributeError('Given storage is not compatible')
+
+        for key, dim in self._cat_dims(self.keys()).items():
+            value1 = self[key]
+            value2 = other[key]
+
+            if key in {'num_nodes', 'num_edges'}:
+                self[key] = value1 + value2
+
+            elif isinstance(value1, list):
+                self[key] = value1 + value2
+
+            elif isinstance(value1, Tensor):
+                self[key] = torch.cat([value1, value2], dim=dim)
+
+            else:
+                raise NotImplementedError(
+                    f"'{self.__class__.__name__}.concat' not yet implemented "
+                    f"for '{type(value1)}'")
+
+        return self
+
+    def is_sorted_by_time(self) -> bool:
+        if 'time' in self:
+            return bool(torch.all(self.time[:-1] <= self.time[1:]))
+        return True
+
+    def sort_by_time(self) -> 'GlobalStorage':
+        if self.is_sorted_by_time():
+            return self
+
+        if 'time' in self:
+            _, perm = torch.sort(self.time, stable=True)
+
+            if self.is_node_attr('time'):
+                keys = self.node_attrs()
+            elif self.is_edge_attr('time'):
+                keys = self.edge_attrs()
+
+            self._select(keys, perm)
+
+        return self
+
+    def snapshot(
+        self,
+        start_time: Union[float, int],
+        end_time: Union[float, int],
+    ) -> 'GlobalStorage':
+        if 'time' in self:
+            mask = (self.time >= start_time) & (self.time <= end_time)
+
+            if self.is_node_attr('time'):
+                keys = self.node_attrs()
+            elif self.is_edge_attr('time'):
+                keys = self.edge_attrs()
+
+            self._select(keys, mask)
+
+            if self.is_node_attr('time') and 'num_nodes' in self:
+                self.num_nodes = int(mask.sum())
+
+        return self
+
+    def up_to(self, time: Union[float, int]) -> 'GlobalStorage':
+        if 'time' in self:
+            return self.snapshot(self.time.min().item(), time)
+        return self
+
 
 class NodeStorage(BaseStorage):
     r"""A storage for node-level information."""


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index 8ebc4e39a..a895c916f 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -517,3 +517,54 @@ def test_data_with_tensor_frame():
     assert data_sub.num_nodes == 3
     for key, value in sub_data.tf.feat_dict.items():
         assert torch.allclose(value, tf.feat_dict[key][mask])
+
+
+@pytest.mark.parametrize('num_nodes', [4])
+@pytest.mark.parametrize('num_edges', [8])
+def test_data_time_handling(num_nodes, num_edges):
+    data = Data(
+        x=torch.randn(num_nodes, 12),
+        edge_index=torch.randint(0, num_nodes, (2, num_edges)),
+        edge_attr=torch.rand((num_edges, 16)),
+        time=torch.arange(num_edges),
+        num_nodes=num_nodes,
+    )
+
+    assert data.is_edge_attr('time')
+    assert not data.is_node_attr('time')
+    assert data.is_sorted_by_time()
+
+    out = data.up_to(5)
+    assert out.num_edges == 6
+    assert torch.allclose(out.x, data.x)
+    assert torch.equal(out.edge_index, data.edge_index[:, :6])
+    assert torch.allclose(out.edge_attr, data.edge_attr[:6])
+    assert torch.equal(out.time, data.time[:6])
+
+    out = data.snapshot(2, 5)
+    assert out.num_edges == 4
+    assert torch.allclose(out.x, data.x)
+    assert torch.equal(out.edge_index, data.edge_index[:, 2:6])
+    assert torch.allclose(out.edge_attr, data.edge_attr[2:6, :])
+    assert torch.equal(out.time, data.time[2:6])
+
+    out = data.sort_by_time()
+    assert data.is_sorted_by_time()
+
+    out = data.concat(data)
+    assert out.num_nodes == 8
+    assert not out.is_sorted_by_time()
+
+    assert torch.allclose(out.x, torch.cat([data.x, data.x], dim=0))
+    assert torch.equal(
+        out.edge_index,
+        torch.cat([data.edge_index, data.edge_index], dim=1),
+    )
+    assert torch.allclose(
+        out.edge_attr,
+        torch.cat([data.edge_attr, data.edge_attr], dim=0),
+    )
+    assert torch.allclose(out.time, torch.cat([data.time, data.time], dim=0))
+
+    out = out.sort_by_time()
+    assert torch.equal(out.time, data.time.repeat_interleave(2))