data.py

"""
Generate synthetic dynamic interaction graph data.
Each time step: a set of edges and node features (spikes).
Labels: binary for each node (or global).
For simplicity, we generate a single binary label per time step (e.g., graph classification).
"""
import numpy as np

def generate_synthetic_data(num_nodes=10, num_steps=10, seed=42):
    rng = np.random.RandomState(seed)
    # Generate a random underlying graph structure (fixed over time?)
    # Use a stochastic block model
    # For simplicity, generate random edges per time step independently.
    adj_mats = []
    node_features = []
    labels = []
    for t in range(num_steps):
        # random graph with probability 0.3
        adj = (rng.rand(num_nodes, num_nodes) < 0.3).astype(np.float64)
        # no self-loops
        np.fill_diagonal(adj, 0)
        adj_mats.append(adj)
        # node features: random binary (1 with prob 0.2)
        x = (rng.rand(num_nodes, 1) < 0.2).astype(np.float64)
        node_features.append(x)
        # label: based on mean node feature (binary threshold)
        label = 1.0 if x.mean() > 0.15 else 0.0
        labels.append(label)
    return adj_mats, node_features, np.array(labels)

def pre_train_data(num_train_steps=7, seed=0):
    return generate_synthetic_data(num_steps=num_train_steps, seed=seed)

def test_data(num_test_steps=3, seed=1):
    return generate_synthetic_data(num_steps=num_test_steps, seed=seed)