Datasets

dg_data

Classes:

• DGData –

  Container for dynamic graph data to be ingested by DGStorage.

DGData dataclass

DGData(
    time_delta: TimeDeltaDG | str,
    time: Tensor,
    edge_mask: Tensor,
    edge_index: Tensor,
    edge_x: Tensor | None = None,
    node_x_mask: Tensor | None = None,
    node_x_nids: Tensor | None = None,
    node_x: Tensor | None = None,
    node_y_mask: Tensor | None = None,
    node_y_nids: Tensor | None = None,
    node_y: Tensor | None = None,
    static_node_x: Tensor | None = None,
    edge_type: Tensor | None = None,
    node_type: Tensor | None = None,
    _split_strategy: SplitStrategy | None = None,
)

Container for dynamic graph data to be ingested by DGStorage.

Stores edge and node events, their timestamps, features, and optional split strategy. Provides methods to split, discretize, and clone the data.

Attributes:

• time_delta (TimeDeltaDG | str) –

  Time granularity of the graph.

• time (Tensor) –

  1D tensor of all event timestamps [num_edge_events + num_node_events + num_node_labels].

• 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_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.

Methods:

• clone –

  Deep copy all tensor and non-tensor fields to create a new DGData object.

• discretize –

  Return a copy of the dataset discretized to a coarser time granularity.

• from_csv –

  Construct a DGData from CSV files containing edge and optional node events.

• from_pandas –

  Construct a DGData from Pandas DataFrames.

• from_raw –

  Construct a DGData from raw tensors for edges, nodes, and features.

• from_tgb –

  Load a DGData from a TGB dataset.

• from_tgb_seq –

  Load a DGData from a TGB-SEQ dataset.

• split –

  Split the dataset according to a strategy.

num_nodes property

num_nodes: int

Global number of nodes in the dataset.

Note: Assumes node ids span a contiguous range, returning max(node_id) + 1.

clone

clone() -> DGData

Deep copy all tensor and non-tensor fields to create a new DGData object.

Source code in tgm/data/dg_data.py

def clone(self) -> DGData:
    """Deep copy all tensor and non-tensor fields to create a new DGData object."""
    cloned_fields = {}
    for f in fields(self):
        val = getattr(self, f.name)
        if isinstance(val, torch.Tensor):
            cloned_fields[f.name] = val.clone()
        else:
            cloned_fields[f.name] = copy.deepcopy(val)

    return replace(self, **cloned_fields)  # type: ignore

discretize

discretize(
    time_delta: TimeDeltaDG | str | None,
    reduce_op: str = 'first',
) -> DGData

Return a copy of the dataset discretized to a coarser time granularity.

Parameters:

• time_delta (TimeDeltaDG | str | None) –

  Target time granularity.

• reduce_op (str, default: 'first' ) –

  Aggregation method for multiple events per bucket. Default 'first'.

Returns:

• DGData ( DGData ) –

  New dataset with discretized timestamps and features.

Raises:

• EventOrderedConversionError –

  If discretization is incompatible with event-ordered granularity

• InvalidDiscretizationError –

  If the target granularity is finer than the current granularity.

Source code in tgm/data/dg_data.py

@log_latency(level=logging.DEBUG)
def discretize(
    self, time_delta: TimeDeltaDG | str | None, reduce_op: str = 'first'
) -> DGData:
    """Return a copy of the dataset discretized to a coarser time granularity.

    Args:
        time_delta (TimeDeltaDG | str | None): Target time granularity.
        reduce_op (str): Aggregation method for multiple events per bucket. Default 'first'.

    Returns:
        DGData: New dataset with discretized timestamps and features.

    Raises:
        EventOrderedConversionError: If discretization is incompatible with event-ordered granularity
        InvalidDiscretizationError: If the target granularity is finer than the current granularity.
    """
    if isinstance(time_delta, str):
        time_delta = TimeDeltaDG(time_delta)
    logger.debug(
        'Discretizing from %s to %s, reduce_op: %s',
        self.time_delta,
        time_delta,
        reduce_op,
    )

    if time_delta is None or self.time_delta == time_delta:
        return self.clone()  # Deepcopy
    if self.time_delta.is_event_ordered or time_delta.is_event_ordered:  # type: ignore
        raise EventOrderedConversionError(
            'Cannot discretize a graph with event-ordered time granularity'
        )
    if self.time_delta.is_coarser_than(time_delta):  # type: ignore
        raise InvalidDiscretizationError(
            f'Cannot discretize to {time_delta} which is strictly'
            f'finer than the self.time_delta: {self.time_delta}'
        )

    valid_ops = ['first']
    if reduce_op not in valid_ops:
        raise ValueError(
            f'Unknown reduce_op: {reduce_op}, expected one of: {valid_ops}'
        )

    # Note: If we simply return a new storage this will 2x our peak memory since the GC
    # won't be able to clean up the current graph storage while `self` is alive.
    time_factor = self.time_delta.convert(time_delta)  # type: ignore
    # Doing time conversion in 64-bit to avoid numerical issues with float cast
    buckets = (self.time.to(torch.float64) * time_factor).floor().int()

    def _get_keep_indices(event_idx: Tensor, ids: Tensor) -> Tensor:
        event_buckets = buckets[event_idx]

        if ids.ndim == 1:
            id_key = ids
        else:
            # Radix-style encoding of edge [src, dst].
            # Collision-free assuming ids >= 0 and no overflow
            src, dst = ids[:, 0], ids[:, 1]
            base = ids.max().item() + 1
            id_key = src * base + dst

        # Radix-style encoding of [bucket, flat_id]
        base = id_key.max().item() + 1
        final_key = event_buckets * base + id_key

        # Stable sort to get adjacent duplicates while preserving original event order
        sort_key, sort_idx = torch.sort(final_key, stable=True)

        # Mark first occurrence in each [bucket, flat_id] group
        is_first = torch.ones_like(sort_key, dtype=torch.bool)
        is_first[1:] = sort_key[1:] != sort_key[:-1]

        # Extract the first [bucket, flat_id] based on our is_first mask
        # Re-sort the final indices so that they index our global timeline chronologically
        keep = sort_idx[is_first]
        keep = keep.sort().values
        return keep

    # Edge events
    edge_mask = _get_keep_indices(self.edge_mask, self.edge_index)
    edge_time = buckets[self.edge_mask][edge_mask]
    edge_index = self.edge_index[edge_mask]
    edge_x = None
    if self.edge_x is not None:
        edge_x = self.edge_x[edge_mask]

    edge_type = None
    if self.edge_type is not None:
        edge_type = self.edge_type[edge_mask]

    # Node events
    node_x_time, node_x_nids, node_x = None, None, None
    if self.node_x_mask is not None:
        node_x_mask = _get_keep_indices(
            self.node_x_mask,
            self.node_x_nids,  # type: ignore
        )
        node_x_time = buckets[self.node_x_mask][node_x_mask]
        node_x_nids = self.node_x_nids[node_x_mask]  # type: ignore
        node_x = None
        if self.node_x is not None:
            node_x = self.node_x[node_x_mask]

    # Node labels
    node_y_time, node_y_nids, node_y = None, None, None
    if self.node_y_mask is not None:
        node_y_mask = _get_keep_indices(
            self.node_y_mask,
            self.node_y_nids,  # type: ignore
        )
        node_y_time = buckets[self.node_y_mask][node_y_mask]
        node_y_nids = self.node_y_nids[node_y_mask]  # type: ignore
        node_y = None
        if self.node_y is not None:
            node_y = self.node_y[node_y_mask]

    # Need a deep copy
    static_node_x = None
    if self.static_node_x is not None:
        logger.debug('Deep copying static node features for coarser DGData')
        static_node_x = self.static_node_x.clone()

    node_type = None
    if self.node_type is not None:  # Need a deep
        logger.debug('Deep copying node_type for coarser DGData')
        node_type = self.node_type.clone()

    return DGData.from_raw(
        time_delta=time_delta,  # new discretized time_delta
        edge_time=edge_time,
        edge_index=edge_index,
        edge_x=edge_x,
        node_x_time=node_x_time,
        node_x_nids=node_x_nids,
        node_x=node_x,
        node_y_time=node_y_time,
        node_y_nids=node_y_nids,
        node_y=node_y,
        static_node_x=static_node_x,
        edge_type=edge_type,
        node_type=node_type,
    )

from_csv classmethod

from_csv(
    edge_file_path: str | Path,
    edge_src_col: str,
    edge_dst_col: str,
    edge_time_col: str,
    edge_x_col: List[str] | None = None,
    node_x_file_path: str | Path | None = None,
    node_x_nids_col: str | None = None,
    node_x_time_col: str | None = None,
    node_x_col: List[str] | None = None,
    node_y_file_path: str | Path | None = None,
    node_y_nids_col: str | None = None,
    node_y_time_col: str | None = None,
    node_y_col: List[str] | None = None,
    static_node_x_file_path: str | Path | None = None,
    static_node_x_col: List[str] | None = None,
    time_delta: TimeDeltaDG | str = 'r',
    edge_type_col: str | None = None,
    node_type_col: str | None = None,
) -> DGData

Construct a DGData from CSV files containing edge and optional node events.

Parameters:

• edge_file_path (str | Path) –

  Path to CSV file containing edges.

• edge_src_col (str) –

  Column name for src nodes.

• edge_dst_col (str) –

  Column name for dst nodes.

• edge_time_col (str) –

  Column name for edge times.

• edge_x_col (List[str] | None, default: None ) –

  Optional edge feature columns.

• node_x_file_path (str | Path | None, default: None ) –

  Optional CSV file for dynamic node features.

• node_x_nids_col (str | None, default: None ) –

  Column name for dynamic node event node ids. Required if node_file_path is specified.

• node_x_time_col (str | None, default: None ) –

  Column name for dynamic node event times. Required if node_file_path is specified.

• node_x_col (List[str] | None, default: None ) –

  Optional dynamic node feature columns.

• node_y_file_path (str | Path | None, default: None ) –

  Optional CSV file for dynamic node labels.

• node_y_nids_col (str | None, default: None ) –

  Column name for dynamic node label node ids.

• node_y_time_col (str | None, default: None ) –

  Column name for dynamic node label times.

• node_y_col (List[str] | None, default: None ) –

  Optional dynamic node label feature columns.

• static_node_x_file_path (str | Path | None, default: None ) –

  Optional CSV file for static node features.

• static_node_x_col (List[str] | None, default: None ) –

  Required if static_node_x_file_path is specified.

• time_delta (TimeDeltaDG | str, default: 'r' ) –

  Time granularity.

• edge_type_col (str | None, default: None ) –

  Column name for edge types.

• node_type_col (str | None, default: None ) –

  Column name for node types.

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.

Source code in tgm/data/dg_data.py

@classmethod
def from_csv(
    cls,
    edge_file_path: str | pathlib.Path,
    edge_src_col: str,
    edge_dst_col: str,
    edge_time_col: str,
    edge_x_col: List[str] | None = None,
    node_x_file_path: str | pathlib.Path | None = None,
    node_x_nids_col: str | None = None,
    node_x_time_col: str | None = None,
    node_x_col: List[str] | None = None,
    node_y_file_path: str | pathlib.Path | None = None,
    node_y_nids_col: str | None = None,
    node_y_time_col: str | None = None,
    node_y_col: List[str] | None = None,
    static_node_x_file_path: str | pathlib.Path | None = None,
    static_node_x_col: List[str] | None = None,
    time_delta: TimeDeltaDG | str = 'r',
    edge_type_col: str | None = None,
    node_type_col: str | None = None,
) -> DGData:
    """Construct a DGData from CSV files containing edge and optional node events.

    Args:
        edge_file_path: Path to CSV file containing edges.
        edge_src_col: Column name for src nodes.
        edge_dst_col: Column name for dst nodes.
        edge_time_col: Column name for edge times.
        edge_x_col: Optional edge feature columns.
        node_x_file_path: Optional CSV file for dynamic node features.
        node_x_nids_col: Column name for dynamic node event node ids. Required if node_file_path is specified.
        node_x_time_col: Column name for dynamic node event times. Required if node_file_path is specified.
        node_x_col: Optional dynamic node feature columns.
        node_y_file_path: Optional CSV file for dynamic node labels.
        node_y_nids_col: Column name for dynamic node label node ids.
        node_y_time_col: Column name for dynamic node label times.
        node_y_col: Optional dynamic node label feature columns.
        static_node_x_file_path: Optional CSV file for static node features.
        static_node_x_col: Required if static_node_x_file_path is specified.
        time_delta: Time granularity.
        edge_type_col: Column name for edge types.
        node_type_col: Column name for node types.

    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.
    """

    def _read_csv(fp: str | pathlib.Path) -> List[dict]:
        # Assumes the whole things fits in memory
        fp = str(fp) if isinstance(fp, pathlib.Path) else fp
        with open(fp, newline='') as f:
            return list(csv.DictReader(f))

    # Read in edge data
    logger.debug('Reading edge_file_path: %s', edge_file_path)
    edge_reader = _read_csv(edge_file_path)
    num_edges = len(edge_reader)

    edge_index = torch.empty((num_edges, 2), dtype=torch.int32)
    timestamps = torch.empty(num_edges, dtype=torch.int64)
    edge_x = None
    if edge_x_col is not None:
        edge_x = torch.empty((num_edges, len(edge_x_col)))
    edge_type = None
    if edge_type_col is not None:
        edge_type = torch.empty(num_edges, dtype=torch.int32)

    for i, row in enumerate(edge_reader):
        edge_index[i, 0] = int(row[edge_src_col])
        edge_index[i, 1] = int(row[edge_dst_col])
        timestamps[i] = int(row[edge_time_col])
        if edge_x_col is not None:
            for j, col in enumerate(edge_x_col):
                edge_x[i, j] = float(row[col])  # type: ignore

        if edge_type_col is not None:
            edge_type[i] = int(row[edge_type_col])  # type: ignore

    # Read in dynamic node data
    node_x_time, node_x_nids, node_x = None, None, None
    if node_x_file_path is not None:
        if node_x_nids_col is None or node_x_time_col is None:
            raise ValueError(
                'specified node_x_file_path without specifying node_x_nids_col and node_x_time_col'
            )
        logger.debug('Reading node_x_file_path: %s', node_x_file_path)
        node_reader = _read_csv(node_x_file_path)
        num_node_events = len(node_reader)

        node_x_time = torch.empty(num_node_events, dtype=torch.int64)
        node_x_nids = torch.empty(num_node_events, dtype=torch.int32)
        if node_x_col is not None:
            node_x = torch.empty((num_node_events, len(node_x_col)))

        for i, row in enumerate(node_reader):
            node_x_time[i] = int(row[node_x_time_col])
            node_x_nids[i] = int(row[node_x_nids_col])
            if node_x_col is not None:
                for j, col in enumerate(node_x_col):
                    node_x[i, j] = float(row[col])  # type: ignore

    # Read in dynamic node labels
    node_y_time, node_y_nids, node_y = None, None, None
    if node_y_file_path is not None:
        if node_y_nids_col is None or node_y_time_col is None:
            raise ValueError(
                'specified node_y_file_path without specifying node_y_nids_col and node_y_time_col'
            )
        logger.debug('Reading node_y_file_path: %s', node_y_file_path)
        node_reader = _read_csv(node_y_file_path)
        num_node_events = len(node_reader)

        node_y_time = torch.empty(num_node_events, dtype=torch.int64)
        node_y_nids = torch.empty(num_node_events, dtype=torch.int32)
        if node_y_col is not None:
            node_y = torch.empty((num_node_events, len(node_y_col)))

        for i, row in enumerate(node_reader):
            node_y_time[i] = int(row[node_y_time_col])
            node_y_nids[i] = int(row[node_y_nids_col])
            if node_y_col is not None:
                for j, col in enumerate(node_y_col):
                    node_y[i, j] = float(row[col])  # type: ignore

    # Read in static node data
    static_node_x = None
    node_type = None
    if static_node_x_file_path is not None:
        if static_node_x_col is None and node_type_col is None:
            raise ValueError(
                'specified static_node_x_file_path without specifying static_node_x_col and node_type_col'
            )
        logger.debug('Reading static_node_x_file_path: %s', static_node_x_file_path)
        static_node_feats_reader = _read_csv(static_node_x_file_path)
        num_nodes = len(static_node_feats_reader)
        if static_node_x_col is not None:
            static_node_x = torch.empty((num_nodes, len(static_node_x_col)))
        if node_type_col is not None:
            node_type = torch.empty(num_nodes, dtype=torch.int32)
        for i, row in enumerate(static_node_feats_reader):
            if static_node_x_col is not None:
                for j, col in enumerate(static_node_x_col):
                    static_node_x[i, j] = float(row[col])  # type: ignore

            if node_type_col is not None:
                node_type[i] = int(row[node_type_col])  # type: ignore

    return cls.from_raw(
        time_delta=time_delta,
        edge_time=timestamps,
        edge_index=edge_index,
        edge_x=edge_x,
        node_x_time=node_x_time,
        node_x_nids=node_x_nids,
        node_x=node_x,
        node_y_time=node_y_time,
        node_y_nids=node_y_nids,
        node_y=node_y,
        static_node_x=static_node_x,
        node_type=node_type,
        edge_type=edge_type,
    )

from_pandas classmethod

from_pandas(
    edge_df: 'pandas.DataFrame',
    edge_src_col: str,
    edge_dst_col: str,
    edge_time_col: str,
    edge_x_col: List[str] | None = None,
    node_x_df: 'pandas.DataFrame' | None = None,
    node_x_nids_col: str | None = None,
    node_x_time_col: str | None = None,
    node_x_col: List[str] | None = None,
    node_y_df: 'pandas.DataFrame' | None = None,
    node_y_nids_col: str | None = None,
    node_y_time_col: str | None = None,
    node_y_col: List[str] | None = None,
    static_node_x_df: 'pandas.DataFrame' | None = None,
    static_node_x_col: List[str] | None = None,
    time_delta: TimeDeltaDG | str = 'r',
    edge_type_col: str | None = None,
    node_type_col: str | None = None,
) -> DGData

Construct a DGData from Pandas DataFrames.

Parameters:

• edge_df ('pandas.DataFrame') –

  DataFrame of edges.

• edge_src_col (str) –

  Column name for src nodes.

• edge_dst_col (str) –

  Column name for dst nodes.

• edge_time_col (str) –

  Column name for edge times.

• edge_x_col (List[str] | None, default: None ) –

  Optional edge feature columns.

• node_x_df ('pandas.DataFrame' | None, default: None ) –

  Optional DataFrame of dynamic node events.

• node_x_nids_col (str | None, default: None ) –

  Column name for dynamic node event node ids. Required if node_file_path is specified.

• node_x_time_col (str | None, default: None ) –

  Column name for dynamic node event times. Required if node_file_path is specified.

• node_x_col (List[str] | None, default: None ) –

  Optional node feature columns.

• node_y_df ('pandas.DataFrame' | None, default: None ) –

  Optional DataFrame of dynamic node labels.

• node_y_nids_col (str | None, default: None ) –

  Column name for dynamic node label node ids.

• node_y_time_col (str | None, default: None ) –

  Column name for dynamic node label times.

• node_y_col (List[str] | None, default: None ) –

  Optional node label feature columns.

• static_node_x_df ('pandas.DataFrame' | None, default: None ) –

  Optional static node features DataFrame.

• static_node_x_col (List[str] | None, default: None ) –

  Required if static_node_x_df is specified.

• time_delta (TimeDeltaDG | str, default: 'r' ) –

  Time granularity.

• edge_type_col (str | None, default: None ) –

  Column name for edge types.

• node_type_col (str | None, default: None ) –

  Column name for node types.

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.

• ImportError –

  If the Pandas package is not resolved in the current python environment.

Source code in tgm/data/dg_data.py

@classmethod
def from_pandas(
    cls,
    edge_df: 'pandas.DataFrame',  # type: ignore
    edge_src_col: str,
    edge_dst_col: str,
    edge_time_col: str,
    edge_x_col: List[str] | None = None,
    node_x_df: 'pandas.DataFrame' | None = None,  # type: ignore
    node_x_nids_col: str | None = None,
    node_x_time_col: str | None = None,
    node_x_col: List[str] | None = None,
    node_y_df: 'pandas.DataFrame' | None = None,  # type: ignore
    node_y_nids_col: str | None = None,
    node_y_time_col: str | None = None,
    node_y_col: List[str] | None = None,
    static_node_x_df: 'pandas.DataFrame' | None = None,  # type: ignore
    static_node_x_col: List[str] | None = None,
    time_delta: TimeDeltaDG | str = 'r',
    edge_type_col: str | None = None,
    node_type_col: str | None = None,
) -> DGData:
    """Construct a DGData from Pandas DataFrames.

    Args:
        edge_df: DataFrame of edges.
        edge_src_col: Column name for src nodes.
        edge_dst_col: Column name for dst nodes.
        edge_time_col: Column name for edge times.
        edge_x_col: Optional edge feature columns.
        node_x_df: Optional DataFrame of dynamic node events.
        node_x_nids_col: Column name for dynamic node event node ids. Required if node_file_path is specified.
        node_x_time_col: Column name for dynamic node event times. Required if node_file_path is specified.
        node_x_col: Optional node feature columns.
        node_y_df: Optional DataFrame of dynamic node labels.
        node_y_nids_col: Column name for dynamic node label node ids.
        node_y_time_col: Column name for dynamic node label times.
        node_y_col: Optional node label feature columns.
        static_node_x_df: Optional static node features DataFrame.
        static_node_x_col: Required if static_node_x_df is specified.
        time_delta: Time granularity.
        edge_type_col: Column name for edge types.
        node_type_col: Column name for node types.

    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.
        ImportError: If the Pandas package is not resolved in the current python environment.
    """

    def _check_pandas_import() -> None:
        try:
            pass
        except ImportError:
            raise ImportError(
                'User requires pandas to initialize a DGraph from a pd.DataFrame'
            )

    _check_pandas_import()

    # Read in edge data
    edge_index = torch.from_numpy(edge_df[[edge_src_col, edge_dst_col]].to_numpy())
    edge_time = torch.from_numpy(edge_df[edge_time_col].to_numpy())
    edge_x = None
    if edge_x_col is not None:
        edge_x = torch.stack([torch.tensor(row) for row in edge_df[edge_x_col]])

    edge_type = None
    if edge_type_col is not None:
        edge_type = torch.from_numpy(edge_df[edge_type_col].to_numpy())

    # Read in dynamic node data
    node_x_time, node_x_nids, node_x = None, None, None
    if node_x_df is not None:
        if node_x_nids_col is None or node_x_time_col is None:
            raise ValueError(
                'specified node_x_df without specifying node_x_nids_col and node_x_time_col'
            )
        node_x_time = torch.as_tensor(node_x_df[node_x_time_col].values)
        node_x_nids = torch.as_tensor(node_x_df[node_x_nids_col].values)
        if node_x_col is not None:
            node_x = torch.stack(
                [torch.tensor(row) for row in node_x_df[node_x_col]]
            )

    # Read in dynamic node labels
    node_y_time, node_y_nids, node_y = None, None, None
    if node_y_df is not None:
        if node_y_nids_col is None or node_y_time_col is None:
            raise ValueError(
                'specified node_y_df without specifying node_y_nids_col and node_y_time_col'
            )
        node_y_time = torch.as_tensor(node_y_df[node_y_time_col].values)
        node_y_nids = torch.as_tensor(node_y_df[node_y_nids_col].values)
        if node_y_col is not None:
            node_y = torch.stack(
                [torch.tensor(row) for row in node_y_df[node_y_col]]
            )

    # Read in static node data
    static_node_x = None
    node_type = None
    if static_node_x_df is not None:
        if static_node_x_col is None and node_type_col is None:
            raise ValueError(
                'specified static_node_x_df without specifying static_node_x_col and node_type'
            )

        if static_node_x_col is not None:
            static_node_x = torch.stack(
                [torch.tensor(row) for row in static_node_x_df[static_node_x_col]]
            )

        if node_type_col is not None:
            node_type = torch.from_numpy(static_node_x_df[node_type_col].to_numpy())

    return cls.from_raw(
        time_delta=time_delta,
        edge_time=edge_time,
        edge_index=edge_index,
        edge_x=edge_x,
        node_x_time=node_x_time,
        node_x_nids=node_x_nids,
        node_x=node_x,
        node_y_time=node_y_time,
        node_y_nids=node_y_nids,
        node_y=node_y,
        static_node_x=static_node_x,
        edge_type=edge_type,
        node_type=node_type,
    )

from_raw classmethod

from_raw(
    edge_time: Tensor,
    edge_index: Tensor,
    edge_x: Tensor | None = None,
    node_x_time: Tensor | None = None,
    node_x_nids: Tensor | None = None,
    node_x: Tensor | None = None,
    node_y_time: Tensor | None = None,
    node_y_nids: Tensor | None = None,
    node_y: Tensor | None = None,
    static_node_x: Tensor | None = None,
    time_delta: TimeDeltaDG | str = 'r',
    edge_type: Tensor | None = None,
    node_type: Tensor | None = None,
) -> DGData

Construct a DGData from raw tensors for edges, nodes, and features.

Automatically combines edge and node timestamps, computes event indices, and validates tensor shapes.

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.

Source code in tgm/data/dg_data.py

@classmethod
def from_raw(
    cls,
    edge_time: Tensor,
    edge_index: Tensor,
    edge_x: Tensor | None = None,
    node_x_time: Tensor | None = None,
    node_x_nids: Tensor | None = None,
    node_x: Tensor | None = None,
    node_y_time: Tensor | None = None,
    node_y_nids: Tensor | None = None,
    node_y: Tensor | None = None,
    static_node_x: Tensor | None = None,
    time_delta: TimeDeltaDG | str = 'r',
    edge_type: Tensor | None = None,
    node_type: Tensor | None = None,
) -> DGData:
    """Construct a DGData from raw tensors for edges, nodes, and features.

    Automatically combines edge and node timestamps, computes event indices,
    and validates tensor shapes.

    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.
    """
    _log_tensor_args(
        edge_time=edge_time,
        edge_index=edge_index,
        edge_x=edge_x,
        node_x_time=node_x_time,
        node_x_nids=node_x_nids,
        node_x=node_x,
        node_y_time=node_y_time,
        node_y_nids=node_y_nids,
        node_y=node_y,
        static_node_x=static_node_x,
        time_delta=time_delta,
        edge_type=edge_type,
        node_type=node_type,
    )
    # Build unified event timeline
    timestamps = edge_time
    event_types = torch.zeros_like(edge_time)
    if node_x_time is not None:
        timestamps = torch.cat([timestamps, node_x_time])
        event_types = torch.cat([event_types, torch.ones_like(node_x_time)])
    if node_y_time is not None:
        timestamps = torch.cat([timestamps, node_y_time])
        event_types = torch.cat(
            [event_types, torch.full_like(node_y_time, fill_value=2)]
        )

    # Compute event masks
    edge_mask = (event_types == 0).nonzero(as_tuple=True)[0]
    node_x_mask = (
        (event_types == 1).nonzero(as_tuple=True)[0]
        if node_x_time is not None
        else None
    )
    node_y_mask = (
        (event_types == 2).nonzero(as_tuple=True)[0]
        if node_y_time is not None
        else None
    )

    return cls(
        time_delta=time_delta,
        time=timestamps,
        edge_mask=edge_mask,
        edge_index=edge_index,
        edge_x=edge_x,
        node_x_mask=node_x_mask,
        node_x_nids=node_x_nids,
        node_x=node_x,
        node_y_mask=node_y_mask,
        node_y_nids=node_y_nids,
        node_y=node_y,
        static_node_x=static_node_x,
        edge_type=edge_type,
        node_type=node_type,
    )

from_tgb classmethod

from_tgb(name: str, **kwargs: Any) -> DGData

Load a DGData from a TGB dataset.

Parameters:

• name (str) –

  Name of the TGB dataset, e.g., 'tgbl-xxxx' or 'tgbn-xxxx'.

• kwargs (Any, default: {} ) –

  Additional dataset-specific arguments.

Returns:

• DGData ( DGData ) –

  Dataset with _split_strategy automatically set to TGBSplit.

Raises:

• ImportError –

  If the TGB package is not resolved in the current python environment.

Notes

• TGBLinkPrediction (tgbl-) and TGBNodePrediction (tgbn-) are supported.
• The split strategy of a TGB dataset cannot be modified.

Source code in tgm/data/dg_data.py

@classmethod
def from_tgb(cls, name: str, **kwargs: Any) -> DGData:
    """Load a DGData from a TGB dataset.

    Args:
        name (str): Name of the TGB dataset, e.g., 'tgbl-xxxx' or 'tgbn-xxxx'.
        kwargs: Additional dataset-specific arguments.

    Returns:
        DGData: Dataset with `_split_strategy` automatically set to `TGBSplit`.

    Raises:
        ImportError: If the TGB package is not resolved in the current python environment.

    Notes:
        - TGBLinkPrediction (`tgbl-`) and TGBNodePrediction (`tgbn-`) are supported.
        - The split strategy of a TGB dataset cannot be modified.
    """
    logger.debug('Loading DGData from TGB dataset: %s', name)
    try:
        from tgb.linkproppred.dataset import LinkPropPredDataset
        from tgb.nodeproppred.dataset import NodePropPredDataset
    except ImportError:
        raise ImportError('TGB required to load TGB data, try `pip install py-tgb`')

    if name.startswith('tgbl-'):
        dataset = LinkPropPredDataset(name=name, **kwargs)
    elif name.startswith('tgbn-'):
        dataset = NodePropPredDataset(name=name, **kwargs)
    elif name.startswith('tkgl-'):
        dataset = LinkPropPredDataset(name=name, **kwargs)
    elif name.startswith('thgl-'):
        dataset = LinkPropPredDataset(name=name, **kwargs)
    else:
        raise ValueError(f'Unknown dataset: {name}')

    data = dataset.full_data

    # IDs are downcast to int32, and features to float32 preventing any runtime warnings.
    # This is safe for all TGB datasets.
    edge_index = torch.stack(
        [
            torch.from_numpy(data['sources']).to(torch.int32),
            torch.from_numpy(data['destinations']).to(torch.int32),
        ],
        dim=1,
    )
    timestamps = torch.from_numpy(data['timestamps']).to(torch.int64)
    if data['edge_feat'] is None:
        edge_x = None
    else:
        edge_x = torch.from_numpy(data['edge_feat']).to(torch.float32)
        if name.startswith('tkgl-'):
            edge_x = torch.cat([edge_x, edge_x])

    node_y_time, node_y_nids, node_y = None, None, None
    if name.startswith('tgbn-'):
        if 'node_label_dict' in data:
            # in TGB, after passing a batch of edges, you find the nearest node event batch in the past
            # in tgbn-trade, validation edge starts at 2010 while the first node event batch starts at 2009.
            # therefore we do (timestamps[0] - 1) to account for this behaviour
            node_label_dict = {
                t: v
                for t, v in data['node_label_dict'].items()
                if (timestamps[0] - 1) <= t < timestamps[-1]
            }
        else:
            raise ValueError(
                f'Failed to construct TGB dataset. Try updating your TGB package: `pip install --upgrade py-tgb`'
            )

        if len(node_label_dict):
            # Node events could be missing from the current data split (e.g. validation)
            num_node_labels, node_label_dim = 0, 0
            for t in node_label_dict:
                for node_id, label in node_label_dict[t].items():
                    num_node_labels += 1
                    node_label_dim = label.shape[0]
            temp_node_y_timestamps = np.zeros(num_node_labels, dtype=np.int64)
            temp_node_y_ids = np.zeros(num_node_labels, dtype=np.int32)
            temp_node_y = np.zeros(
                (num_node_labels, node_label_dim), dtype=np.float32
            )
            idx = 0
            for t in node_label_dict:
                for node_id, label in node_label_dict[t].items():
                    temp_node_y_timestamps[idx] = t
                    temp_node_y_ids[idx] = node_id
                    temp_node_y[idx] = label
                    idx += 1
            node_y_time = torch.from_numpy(temp_node_y_timestamps)
            node_y_nids = torch.from_numpy(temp_node_y_ids)
            node_y = torch.from_numpy(temp_node_y)

    # Read static node features if they exist
    static_node_x = None
    if dataset.node_feat is not None:
        static_node_x = torch.from_numpy(dataset.node_feat).to(torch.float32)

    edge_type = None
    node_type = None
    if name.startswith('thgl'):
        if 'edge_type' not in data or not hasattr(dataset, 'node_type'):
            raise ValueError(
                f'Failed to construct TGB dataset. Try updating your TGB package: `pip install --upgrade py-tgb`'
            )
        edge_type = torch.from_numpy(data['edge_type']).to(torch.int32)
        node_type = torch.from_numpy(dataset.node_type).to(torch.int32)

    if name.startswith('tkgl'):
        if 'edge_type' not in data:
            raise ValueError(
                f'Failed to construct TGB dataset. Try updating your TGB package: `pip install --upgrade py-tgb`'
            )
        edge_type = torch.from_numpy(data['edge_type']).to(torch.int32)

    raw_times = torch.from_numpy(data['timestamps'])
    split_bounds = {}
    for split_name, mask in {
        'train': dataset.train_mask,
        'val': dataset.val_mask,
        'test': dataset.test_mask,
    }.items():
        times = raw_times[mask]
        split_bounds[split_name] = (int(times.min()), int(times.max()))

    data = cls.from_raw(
        time_delta=TGB_TIME_DELTAS[name],
        edge_time=timestamps,
        edge_index=edge_index,
        edge_x=edge_x,
        node_y_time=node_y_time,
        node_y_nids=node_y_nids,
        node_y=node_y,
        static_node_x=static_node_x,
        edge_type=edge_type,
        node_type=node_type,
    )

    data._split_strategy = TGBSplit(split_bounds)
    logger.debug('Finished loading DGData from TGB dataset: %s', name)
    return data

from_tgb_seq classmethod

from_tgb_seq(name: str, **kwargs: Any) -> DGData

Load a DGData from a TGB-SEQ dataset.

Parameters:

• name (str) –

  Name of the TGB-SEQ dataset.

• kwargs (Any, default: {} ) –

  Additional dataset-specific arguments.

Returns:

• DGData ( DGData ) –

  Dataset with _split_strategy automatically set to TGBSplit.

Raises:

• ImportError –

  If the TGB-SEQ package is not resolved in the current python environment.

Notes

• The split strategy of a TGB-SEQ dataset cannot be modified.
• If a data root location is not specified, we default store location to './data'.

Source code in tgm/data/dg_data.py

@classmethod
def from_tgb_seq(cls, name: str, **kwargs: Any) -> DGData:
    """Load a DGData from a TGB-SEQ dataset.

    Args:
        name (str): Name of the TGB-SEQ dataset.
        kwargs: Additional dataset-specific arguments.

    Returns:
        DGData: Dataset with `_split_strategy` automatically set to `TGBSplit`.

    Raises:
        ImportError: If the TGB-SEQ package is not resolved in the current python environment.

    Notes:
        - The split strategy of a TGB-SEQ dataset cannot be modified.
        - If a data root location is not specified, we default store location to './data'.
    """
    logger.debug('Loading DGData from TGB-SEQ dataset: %s', name)
    try:
        from tgb_seq.LinkPred.dataloader import TGBSeqLoader
    except ImportError:
        raise ImportError(
            'TGB-SEQ required to load TGB data, try `pip install tgb-seq`'
        )

    if 'root' not in kwargs:
        kwargs['root'] = './data'
    data = TGBSeqLoader(name=name, **kwargs)

    # IDs are downcast to int32, and features to float32 preventing any runtime warnings.
    # This is safe for all TGB datasets.
    edge_index = torch.stack(
        [
            torch.from_numpy(data.src_node_ids).to(torch.int32),
            torch.from_numpy(data.dst_node_ids).to(torch.int32),
        ],
        dim=1,
    )

    timestamps = torch.from_numpy(data.node_interact_times).to(torch.int64)
    edge_x = None
    if data.edge_features is not None:
        edge_x = torch.from_numpy(data.edge_features).to(torch.float32)

    # Read static node features if they exist
    static_node_x = None
    if data.node_features is not None:
        static_node_x = torch.from_numpy(data.node_features).to(torch.float32)

    split_bounds = {}
    for split_name, mask in {
        'train': data.train_mask,
        'val': data.val_mask,
        'test': data.test_mask,
    }.items():
        times = data.node_interact_times[mask]
        split_bounds[split_name] = (int(times.min()), int(times.max()))

    data = cls.from_raw(
        time_delta=TGB_SEQ_TIME_DELTAS[name],
        edge_time=timestamps,
        edge_index=edge_index,
        edge_x=edge_x,
        static_node_x=static_node_x,
    )

    data._split_strategy = TGBSplit(split_bounds)
    logger.debug('Finished loading DGData from TGB-SEQ dataset: %s', name)
    return data

split

split(
    strategy: SplitStrategy | None = None,
) -> Tuple[DGData, ...]

Split the dataset according to a strategy.

Parameters:

• strategy (SplitStrategy | None, default: None ) –

  Optional strategy to override the default. If None, uses _split_strategy or defaults to TemporalRatioSplit.

Returns:

• Tuple[DGData, ...] –

  Tuple[DGData, ...]: Split datasets (train/val/test).

Raises:

• ValueError –

  If attempting to override the split strategy for TGB datasets.

Notes

• Splits preserve the underlying storage; only indices are filtered.

Source code in tgm/data/dg_data.py

def split(self, strategy: SplitStrategy | None = None) -> Tuple[DGData, ...]:
    """Split the dataset according to a strategy.

    Args:
        strategy (SplitStrategy | None): Optional strategy to override the
            default. If None, uses `_split_strategy` or defaults to `TemporalRatioSplit`.

    Returns:
        Tuple[DGData, ...]: Split datasets (train/val/test).

    Raises:
        ValueError: If attempting to override the split strategy for TGB datasets.

    Notes:
        - Splits preserve the underlying storage; only indices are filtered.
    """
    strategy = strategy or self._split_strategy or TemporalRatioSplit()

    if (
        isinstance(self._split_strategy, TGBSplit)
        and strategy is not self._split_strategy
    ):
        raise ValueError('Cannot override split strategy for TGB datasets')

    return strategy.apply(self)

Dataset Splitting

SplitStrategy

Bases: ABC

Abstract base class for splitting temporal graph datasets.

Implementations of this class define the logic for dividing a DGData object into one or more subsets (train/val/test) based on temporal information.

Methods:

• apply –

  Split the dataset and return one or more subsets.

apply abstractmethod

apply(data: 'DGData') -> Tuple['DGData', ...]

Split the dataset and return one or more subsets.

Parameters:

• data (DGData) –

  The temporal graph dataset to split.

Returns:

• Tuple['DGData', ...] –

  Tuple[DGData, ...]: Split datasets according to the strategy.

Source code in tgm/data/split.py

@abstractmethod
def apply(self, data: 'DGData') -> Tuple['DGData', ...]:  # type: ignore
    r"""Split the dataset and return one or more subsets.

    Args:
        data (DGData): The temporal graph dataset to split.

    Returns:
        Tuple[DGData, ...]: Split datasets according to the strategy.
    """

TemporalSplit dataclass

TemporalSplit(val_time: int, test_time: int)

Bases: SplitStrategy

Split a temporal graph dataset based on absolute timestamp boundaries.

Parameters:

• val_time (int) –

  The timestamp separating training and validation data.

• test_time (int) –

  The timestamp separating validation and test data.

Raises:

• ValueError –

  If not 0 <= val_time <= test_time.

Note

This strategy assigns edges and nodes to splits based on whether their timestamps fall within the corresponding intervals: - Train: (-inf, val_time) - Validation: [val_time, test_time) - Test: [test_time, inf)

TemporalRatioSplit dataclass

TemporalRatioSplit(
    train_ratio: float = 0.7,
    val_ratio: float = 0.15,
    test_ratio: float = 0.15,
)

Bases: SplitStrategy

Split a temporal graph dataset according to relative ratios of time.

Parameters:

• train_ratio (float, default: 0.7 ) –

  Fraction of data to assign to training. Defaults to 0.7.

• val_ratio (float, default: 0.15 ) –

  Fraction of data to assign to validation. Defaults to 0.15.

• test_ratio (float, default: 0.15 ) –

  Fraction of data to assign to test. Defaults to 0.15.

Raises:

• ValueError –

  If any ratio is negative.

• ValueError –

  If the ratios do not sum to 1.0 within tolerance.

Note

The dataset timestamps are assumed to be sorted. The ratios are applied cumulatively to the total temporal span to determine absolute split boundaries.

TGBSplit dataclass

TGBSplit(split_bounds: Dict[str, Tuple[int, int]])

Bases: SplitStrategy

Split a temporal graph dataset using pre-specified edge time bounds.

Parameters:

• split_bounds (Dict[str, Tuple[int, int]]) –

  Mapping from split names ('train', 'val', 'test') to (min_time, max_time) intervals for edges.

Note

Nodes are included in a split if their timestamps fall within the corresponding edge interval (or slightly before the min_time of edges).