Skip to content

Graph Data Structures

k3_node.data provides graph data structures, datasets, and storage abstractions compatible with NumPy, PyTorch, TensorFlow, and JAX tensors.

Graph Containers

Data

Data represents a single homogeneous graph, holding node attributes, edge indices, edge attributes, and graph-level properties.

import numpy as np
from k3_node.data import Data

# Graph with 3 nodes and 4 directed edges
x = np.array([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]], dtype=np.float32)
edge_index = np.array([[0, 1, 1, 2],
                       [1, 0, 2, 1]], dtype=np.int64)

data = Data(x=x, edge_index=edge_index)
print(data.num_nodes)  # 3
print(data.num_edges)  # 4

k3_node.data.Data

Bases: BaseData

A data object describing a homogeneous graph.

inputs property

Returns the tuple of input tensors (x, edge_index) or (x, edge_index, edge_attr).

accuracy(logits_or_pred, mask='test_mask')

Convenience method to calculate classification accuracy.

Parameters:

Name Type Description Default
logits_or_pred Any

Model prediction output (either class logits or predicted labels).

 required
mask str

The mask attribute (e.g. 'test_mask') to evaluate on. If :obj:None, evaluates across all nodes. (default: 'test_mask')

 'test_mask'

edge_subgraph(subset)

Returns the graph with only the edges in subset (a boolean edge mask or edge indices). All nodes are kept; every edge-level attribute is filtered.

Example 
import numpy as np
from k3_node.data import Data

data = Data(edge_index=np.array([[0, 1, 2], [1, 2, 0]]), edge_type=np.array([0, 1, 0]), num_nodes=3)
train = data.edge_subgraph(np.array([True, False, True]))
print(tuple(train.edge_index.shape), train.num_nodes)  # (2, 2) 3

subgraph(subset)

Returns the induced subgraph for subset nodes.

to_generator(mask='train_mask', repeat=True)

Generates tuples of ((x, edge_index), y, mask) or ((x, edge_index, edge_attr), y, mask) ready for training directly with Keras model.fit().

Parameters:

Name Type Description Default
mask str

The name of the mask attribute (e.g. 'train_mask', 'val_mask', 'test_mask') to use as sample_weight for loss masking. If :obj:None, no mask is applied. (default: 'train_mask')

 'train_mask'
repeat bool

Whether to yield infinitely for Keras generator training. (default: :obj:True)

 True

HeteroData

HeteroData represents heterogeneous graphs with multiple node types and edge types.

from k3_node.data import HeteroData

data = HeteroData()
data['paper'].x = paper_features
data['author'].x = author_features
data['author', 'writes', 'paper'].edge_index = author_writes_paper_edges

print(data.node_types)  # ['paper', 'author']
print(data.edge_types)  # [('author', 'writes', 'paper')]

k3_node.data.HeteroData

Bases: BaseData

A data object describing a heterogeneous graph.

get_edge_store(src, rel, dst)

Gets the EdgeStorage object of a particular edge type given by (src, rel, dst).

get_node_store(key)

Gets the NodeStorage object of a particular node type.

Batch & HeteroBatch

Batch combines multiple Data objects into a single giant disjoint graph with diagonal adjacency block structure and a batch vector tracking graph membership.

from k3_node.data import Batch, Data

graph1 = Data(x=x1, edge_index=edge_index1)
graph2 = Data(x=x2, edge_index=edge_index2)

batch = Batch.from_data_list([graph1, graph2])
print(batch.batch)       # [0, 0, 0, ..., 1, 1, 1]
print(batch.num_graphs)  # 2

# Deconstruct back into individual graphs
data_list = batch.to_data_list()

k3_node.data.Batch

Bases: Data

A data object describing a batch of graphs as one big (disconnected) graph.

TemporalData & HypergraphData

  • TemporalData: Encapsulates continuous-time dynamic interaction streams with timestamps (t), source nodes (src), destination nodes (dst), and edge attributes (msg).
  • HypergraphData: Encapsulates hypergraphs where hyperedges can connect an arbitrary number of nodes.

k3_node.data.TemporalData

Bases: BaseData

A data object composed of a stream of events describing a temporal graph.

train_val_test_split(val_ratio=0.15, test_ratio=0.15)

Splits the events chronologically into training, validation and test events (the last val_ratio + test_ratio of the time span go to validation and test).

k3_node.data.HypergraphData

Bases: Data

A data object describing a hypergraph.

Datasets

Dataset & InMemoryDataset

  • Dataset: Base class for graph datasets, supporting raw file downloading, processing, and indexing.
  • InMemoryDataset: Dataset that fits completely into memory, with fast serialization/deserialization.
  • OnDiskDataset: Out-of-core dataset backed by key-value stores (Database, SQLiteDatabase, RocksDatabase).
from k3_node.data import InMemoryDataset

class MyGraphDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
        self.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ['data.csv']

    @property
    def processed_file_names(self):
        return ['data.pt']

    def process(self):
        # Process raw files and save graphs
        ...

k3_node.data.InMemoryDataset

Bases: Dataset

Dataset base class for in-memory graph collections.

Remote & Large-Scale Stores

  • FeatureStore: Standard interface for external or out-of-memory node/edge feature storage.
  • GraphStore: Standard interface for graph topology and edge indexing.
  • SQLiteDatabase: SQLite3-backed persistent graph and attribute storage.
  • RocksDatabase: RocksDB-backed key-value storage.

Download & Extraction Utilities

  • download_url(url, folder, filename=None): Downloads a file over HTTP/HTTPS with progress bar.
  • download_google_url(id, folder, filename): Downloads large pre-trained checkpoints or datasets directly from Google Drive.
  • extract_zip(path, folder) / extract_tar(path, folder) / extract_gz(path, folder) / extract_bz2(path, folder): Archive extraction helpers.
  • makedirs(path): Recursive directory creation helper.