can't split HeteroData with RandomNodeSplit

[Rexedoziem]

import networkx as nx
import pandas as pd
from collections import defaultdict
import torch
from torch_geometric.data import Data, HeteroData
import gensim.downloader as api
from gensim.models import Word2Vec
from sklearn.preprocessing import StandardScaler, OneHotEncoder
from torch_geometric.loader import NeighborLoader
import torch_geometric.transforms as T
from torch_geometric.transforms import RandomNodeSplit

Load pre-trained Word2Vec model

wv_model = api.load("glove-wiki-gigaword-50")

def encode_text(text):
"""Encode text using Word2Vec or GloVe"""
words = text.split()
embeddings = []
for word in words:
try:
embedding = wv_model.get_vector(word)
embeddings.append(embedding)
except KeyError:
# If the word is not in the vocabulary, skip it
pass
if len(embeddings) == 0:
# If no words are found, return a random vector
return torch.randn(wv_model.vector_size)
return torch.mean(torch.tensor(embeddings), dim=0)

scaler = StandardScaler()
encoder = OneHotEncoder(sparse=False)

def preprocess_data(data):
"""Preprocess the data by handling missing values, splitting categories, and encoding features."""
data = data.fillna({'text': '', 'categories': ''})
data['categories'] = data['categories'].str.split(', ')
# Scale numerical features
data = scale_numerical_features(data, ['longitude', 'latitude', 'review_count_x', 'review_count'])
# Encode categorical features
data = encode_categorical_features(data, ['city'])
return data

def scale_numerical_features(data, features):
"""Scale numerical features using StandardScaler"""
scaled_data = data.copy()
scaled_data[features] = scaler.fit_transform(scaled_data[features])
return scaled_data

def encode_categorical_features(data, features):
"""Encode categorical features using OneHotEncoder"""
encoded_data = data.copy()
encoded_data[features] = encoder.fit_transform(encoded_data[features])
return encoded_data

def create_bipartite_graphs(data):
"""Create user-restaurant and user-category bipartite graphs."""
user_restaurant_graph = nx.Graph()
user_category_graph = nx.Graph()

# Unique nodes
unique_users = data['user_id'].unique()
unique_restaurants = data['business_id'].unique()
unique_categories = set(category for categories in data['categories'] for category in categories)

# Add nodes to user-restaurant graph
user_restaurant_graph.add_nodes_from(unique_users, bipartite='user')
user_restaurant_graph.add_nodes_from(unique_restaurants, bipartite='restaurant')

# Add nodes to user-category graph
user_category_graph.add_nodes_from(unique_users, bipartite='user')
user_category_graph.add_nodes_from(unique_categories, bipartite='category')

# Add edges to user-restaurant graph
user_restaurant_edges = [(row['user_id'], row['business_id']) for _, row in data.iterrows()]
user_restaurant_graph.add_edges_from(user_restaurant_edges)

# Add 'stars_x' as a weighted edge attribute to user-restaurant graph
star_x_edges = [(row['user_id'], row['business_id'], {'star_x': row['stars_x']}) for _, row in data.iterrows()]
user_restaurant_graph.add_edges_from(star_x_edges)

# Add edges to user-category graph
user_category_edges = [(row['user_id'], category) for _, row in data.iterrows() for category in row['categories']]
user_category_graph.add_edges_from(user_category_edges)

# Add 'stars_y' as a weighted edge attribute to user-category graph
star_y_edges = [(row['user_id'], category, {'star_y': row['stars_y']}) for _, row in data.iterrows() for category in row['categories']]
user_category_graph.add_edges_from(star_y_edges)

return user_restaurant_graph, user_category_graph

def add_node_features(user_restaurant_graph, user_category_graph, data):
"""Add node features to the bipartite graphs."""
# Add 'review_count' as a node feature
user_review_count = dict(zip(data['user_id'], data['review_count_x']))
restaurant_review_count = dict(zip(data['business_id'], data['review_count']))
nx.set_node_attributes(user_restaurant_graph, user_review_count, name='review_count_user')
nx.set_node_attributes(user_restaurant_graph, restaurant_review_count, name='review_count_restaurant')

# Add 'city' as a node feature
restaurant_city = dict(zip(data['business_id'], data['city']))
nx.set_node_attributes(user_restaurant_graph, restaurant_city, name='city_restaurant')

# Add 'longitude' and 'latitude' as node features for restaurants
restaurant_longitude = dict(zip(data['business_id'], data['longitude']))
restaurant_latitude = dict(zip(data['business_id'], data['latitude']))
nx.set_node_attributes(user_restaurant_graph, restaurant_longitude, name='longitude_restaurant')
nx.set_node_attributes(user_restaurant_graph, restaurant_latitude, name='latitude_restaurant')

# Add categories as a node feature for users
user_categories = defaultdict(list)
for _, row in data.iterrows():
    user_categories[row['user_id']].extend(row['categories'])
nx.set_node_attributes(user_category_graph, {user: list(set(categories)) for user, categories in user_categories.items()}, name='categories_user')

def add_edge_features(user_restaurant_graph, user_category_graph, data):
"""Add edge features to the bipartite graphs."""
# Add 'text' as an edge feature
user_restaurant_text = {(row['user_id'], row['business_id']): row['text'] for _, row in data.iterrows()}
user_category_text = {(row['user_id'], category): row['text'] for _, row in data.iterrows() for category in row['categories']}
nx.set_edge_attributes(user_restaurant_graph, user_restaurant_text, name='review_text')
nx.set_edge_attributes(user_category_graph, user_category_text, name='review_text')

# Add 'name' and 'categories' as edge features for user-restaurant graph
user_restaurant_name = {(row['user_id'], row['business_id']): row['name'] for _, row in data.iterrows()}
user_restaurant_categories = {(row['user_id'], row['business_id']): row['categories'] for _, row in data.iterrows()}
nx.set_edge_attributes(user_restaurant_graph, user_restaurant_name, name='restaurant_name')
nx.set_edge_attributes(user_restaurant_graph, user_restaurant_categories, name='restaurant_categories')
# Add 'star_x' as an edge feature with weights
star_x_edges = [(row['user_id'], row['business_id'], {'star_x': row['stars_x']}) for _, row in data.iterrows()]
user_restaurant_graph.add_edges_from(star_x_edges)
# Add 'stars_y' as a weighted edge attribute to user-category graph
star_y_edges = [(row['user_id'], category, {'star_y': row['stars_y']}) for _, row in data.iterrows() for category in row['categories']]
user_category_graph.add_edges_from(star_y_edges)

def encode_node_features(node_data):
"""Encode node features using node attributes and review text."""
#print(f"Node data: {node_data}") # Add this line
node_attributes = []

# Encode categorical attributes
city = node_data.get('city', '')
if city:
    city_encoding = encoder.transform([[city]]).toarray().flatten()
    node_attributes.extend(city_encoding)
else:
    node_attributes.extend([0] * len(encoder.categories_[0]))

# Encode numerical attributes
review_count = node_data.get('review_count', 0)
node_attributes.append(review_count)

# Encode boolean attributes
is_open = node_data.get('is_open', 0)
node_attributes.append(float(is_open))

# Encode review text
review_text = encode_text(node_data.get('review_text', ''))
node_attributes.extend(review_text.tolist())

return torch.tensor(node_attributes)

def create_pyg_data(graph, node_type1, node_type2, node_type3=None):
"""Create a PyG Data or HeteroData object from a bipartite graph."""
data = HeteroData()

# Add node features for node_type1
node_type1_features = []
for _, node_data in graph.nodes(data=True):
    if node_data.get('bipartite') == node_type1:
        node_type1_features.append(encode_node_features(node_data))
data[node_type1].x = torch.stack(node_type1_features)

# Add node features for node_type2
node_type2_features = []
for _, node_data in graph.nodes(data=True):
    if node_data.get('bipartite') == node_type2:
        node_type2_features.append(encode_node_features(node_data))
data[node_type2].x = torch.stack(node_type2_features)

# Add node features for node_type3 (if provided)
if node_type3 is not None:
    node_type3_features = []
    for _, node_data in graph.nodes(data=True):
        if node_data.get('bipartite') == node_type3:
            node_type3_features.append(encode_node_features(node_data))
    data[node_type3].x = torch.stack(node_type3_features)

# Add edge indices
node_type_mapping = {'user': 0, 'restaurant': 1, 'category': 2}
edges = [(node_type_mapping[graph.nodes[u]['bipartite']], node_type_mapping[graph.nodes[v]['bipartite']]) for u, v in graph.edges()]
row, col = zip(*edges)
data[(node_type1, 'to', node_type2)].edge_index = torch.tensor([row, col], dtype=torch.long)
if node_type3 is not None:
    data[(node_type1, 'to', node_type3)].edge_index = torch.tensor([row, col], dtype=torch.long)

# Add edge features
edge_features = [encode_text(graph.edges[u, v]['review_text']) for u, v in graph.edges()]
data[(node_type1, 'to', node_type2)].edge_attr = torch.stack(edge_features)
if node_type3 is not None:
    data[(node_type1, 'to', node_type3)].edge_attr = torch.stack(edge_features)

# Add 'star_x' as an edge feature for user-restaurant graph
user_restaurant_star_x = [graph.edges[u, v].get('star_x', 0) for u, v in graph.edges() if graph.nodes[u]['bipartite'] == node_type1 and graph.nodes[v]['bipartite'] == node_type2]
data[(node_type1, 'to', node_type2)].edge_weight = torch.tensor(user_restaurant_star_x)

# Add 'star_y' as an edge feature for user-category graph
if node_type3 is not None:
    user_category_star_y = [graph.edges[u, v].get('star_y', 0) for u, v in graph.edges() if graph.nodes[u]['bipartite'] == node_type1 and graph.nodes[v]['bipartite'] == node_type3]
    data[(node_type1, 'to', node_type3)].edge_weight = torch.tensor(user_category_star_y)

return data

def create_combined_data(user_restaurant_graph, user_category_graph, preprocessed_data):
"""Create the combined HeteroData object."""
# Add node features
add_node_features(user_restaurant_graph, user_category_graph, preprocessed_data)

# Add edge features
add_edge_features(user_restaurant_graph, user_category_graph, preprocessed_data)

# Create PyG Data objects
user_restaurant_data = create_pyg_data(user_restaurant_graph, 'user', 'restaurant')
user_category_data = create_pyg_data(user_category_graph, 'user', 'category')

# Create PyG homogeneous Data object (optional)
combined_graph = nx.compose(user_restaurant_graph, user_category_graph)
combined_data = create_pyg_data(combined_graph, 'user', 'restaurant', 'category')
# Split the combined_data into train, validation, and test sets based on 'user' node type
combined_data = RandomNodeSplit(split='train_rest', num_val=0.1, num_test=0.1)(combined_data)
combined_data = T.ToUndirected()(combined_data)
print(f'data: {combined_data}')

return combined_data, combined_graph

def main():
# Load the dataset
data = pd.read_csv("/kaggle/working/subset_review_merged1.csv")

# Preprocess the data
preprocessed_data = preprocess_data(data)
# Create bipartite graphs
user_restaurant_graph, user_category_graph = create_bipartite_graphs(preprocessed_data)

# Create the combined HeteroData object and the combined graph
combined_data, combined_graph = create_combined_data(user_restaurant_graph, user_category_graph, preprocessed_data)

return combined_graph, combined_data

# Define the is_connected function
def is_connected(graph):
    """Check if a graph is connected using BFS."""
    # Initialize visited set and queue for BFS
    visited = set()
    queue = []

    # Choose an arbitrary starting node
    start_node = next(iter(graph.nodes()))

    # Start BFS from the chosen node
    queue.append(start_node)
    visited.add(start_node)

    while queue:
        # Dequeue a node from the queue
        node = queue.pop(0)

        # Visit neighbors of the current node
        for neighbor in graph.neighbors(node):
            if neighbor not in visited:
                visited.add(neighbor)
                queue.append(neighbor)

    # If all nodes are visited, the graph is connected
    return len(visited) == len(graph.nodes())

# Check if the combined graph is connected
combined_connected = is_connected(combined_graph)
print("Is the combined graph connected?", combined_connected)

# Check if the combined graph is bipartite
combined_bipartite = nx.is_bipartite(combined_graph)

# Visualization
plt.figure(figsize=(10, 6))
nx.draw(combined_graph, with_labels=True)
plt.title(f'Combined Bipartite Graph\nBipartite: {combined_bipartite}')
plt.show()

return combined_graph, combined_data

if name == 'main':
combined_graph, combined_data = main()

# Access edge attributes for ('user', 'to', 'restaurant') edge type
user_restaurant_edge_index = combined_data[('user', 'to', 'restaurant')].edge_index
user_restaurant_edge_attr = combined_data[('user', 'to', 'restaurant')].edge_attr
user_restaurant_edge_weight = combined_data[('user', 'to', 'restaurant')].edge_weight

# Access edge attributes for ('user', 'to', 'category') edge type
user_category_edge_index = combined_data[('user', 'to', 'category')].edge_index
user_category_edge_attr = combined_data[('user', 'to', 'category')].edge_attr
user_category_edge_weight = combined_data[('user', 'to', 'category')].edge_weight
# Print the dimensions of edge attributes for each edge type
print(f"Shape of data['user'].x: {combined_data['user'].x.shape}")

print("User to Restaurant Edge Index shape:", user_restaurant_edge_index.shape)
print("User to Restaurant Edge Attribute shape:", user_restaurant_edge_attr.shape)

This is the combined_data even after the split:
data: HeteroData(
user={ x=[583, 53] },
restaurant={ x=[142, 53] },
category={ x=[122, 53] },
(user, to, restaurant)={
edge_index=[2, 3669],
edge_attr=[3669, 50],
edge_weight=[592]
},
(user, to, category)={
edge_index=[2, 3669],
edge_attr=[3669, 50],
edge_weight=[3077]
},
(restaurant, rev_to, user)={
edge_index=[2, 3669],
edge_attr=[3669, 50]
},
(category, rev_to, user)={
edge_index=[2, 3669],
edge_attr=[3669, 50]
}
)

My issues is why after using the RandomNodeSplit, it didn't split the data to train, val and test_masks.. So is there anything i'm doing wrong?

[rusty1s]

Where do you define your labels? Currently, RandomNodeSplit will split based on the node types that define a label (y by default).

[Rexedoziem]

star_x is my label, but was not defined, so i ended up splitting based on the node type using the RandomNodeSplit key of 'x' instead of y for my analysis. so i dont know if it can work too. This is what i got as the data:
combined data: HeteroData(
user={
num_nodes=583,
x=[583, 53],
train_mask=[583],
val_mask=[583],
test_mask=[583]
},
restaurant={
num_nodes=142,
x=[142, 53],
train_mask=[142],
val_mask=[142],
test_mask=[142]
},
category={
num_nodes=122,
x=[122, 53],
train_mask=[122],
val_mask=[122],
test_mask=[122]
},
(user, to, restaurant)={
edge_index=[2, 3669],
edge_attr=[3669, 50],
edge_weight=[592]
},
(user, to, category)={
edge_index=[2, 3669],
edge_attr=[3669, 50],
edge_weight=[3077]
},
(restaurant, rev_to, user)={
edge_index=[2, 3669],
edge_attr=[3669, 50]
},
(category, rev_to, user)={
edge_index=[2, 3669],
edge_attr=[3669, 50]
}
)

[rusty1s]

Yes, this looks better.

[Rexedoziem]

Thank you so much, I did everything as expected.

[Rexedoziem]

combined data: HeteroData(
user={
num_nodes=583,
x=[583, 53],
y=[583]
train_mask=[583],
val_mask=[583],
test_mask=[583]
},
restaurant={
num_nodes=142,
x=[142, 53],
},
category={
num_nodes=122,
x=[122, 53],
},
(user, to, restaurant)={
edge_index=[2, 3669],
edge_attr=[3669, 50],
},
(user, to, category)={
edge_index=[2, 3669],
edge_attr=[3669, 50],
edge_weight=[3077]
},
(restaurant, rev_to, user)={
edge_index=[2, 3669],
edge_attr=[3669, 50]
},
(category, rev_to, user)={
edge_index=[2, 3669],
edge_attr=[3669, 50]
}
)

I got this after putting my Label(y) in place… Is it correct?

[rusty1s]

Sorry for late reply :( I was actually referring to the processed data, so I don't need to run your preprocess routine. E.g., via

torch.save(data, 'data.pt')

[Rexedoziem]

Okay now I get it, I tried sending it here on GitHub, but it wasn’t going through because it doesn’t support the .pt file.

[Rexedoziem]

So do I send it across through another platform or?

[rusty1s]

you can upload it anywhere or send it to me on Slack.

[Rexedoziem]

It’s on slack

[Rexedoziem]

Yes, because the code is centered on the user node.