Can't split HeteroData with RandomNodeSplit

username: Rexedoziem
cooked: <p>import networkx as nx<br>
import pandas as pd<br>
from collections import defaultdict<br>
import torch<br>
from torch_geometric.data import Data, HeteroData<br>
import gensim.downloader as api<br>
from gensim.models import Word2Vec<br>
from sklearn.preprocessing import StandardScaler, OneHotEncoder<br>
from torch_geometric.loader import NeighborLoader<br>
import torch_geometric.transforms as T<br>
from torch_geometric.transforms import RandomNodeSplit</p>
<h1><a name="load-pre-trained-word2vec-model-1" class="anchor" href="#load-pre-trained-word2vec-model-1"></a>Load pre-trained Word2Vec model</h1>
<p>wv_model = api.load(“glove-wiki-gigaword-50”)</p>
<p>def encode_text(text):<br>
“”“Encode text using Word2Vec or GloVe”“”<br>
words = text.split()<br>
embeddings = <span class="chcklst-box fa fa-square-o fa-fw"></span><br>
for word in words:<br>
try:<br>
embedding = wv_model.get_vector(word)<br>
embeddings.append(embedding)<br>
except KeyError:<br>
# If the word is not in the vocabulary, skip it<br>
pass<br>
if len(embeddings) == 0:<br>
# If no words are found, return a random vector<br>
return torch.randn(wv_model.vector_size)<br>
return torch.mean(torch.tensor(embeddings), dim=0)</p>
<p>scaler = StandardScaler()<br>
encoder = OneHotEncoder(sparse=False)</p>
<p>def preprocess_data(data):<br>
“”“Preprocess the data by handling missing values, splitting categories, and encoding features.”“”<br>
data = data.fillna({‘text’: ‘’, ‘categories’: ‘’})<br>
data[‘categories’] = data[‘categories’].str.split(', ')<br>
# Scale numerical features<br>
data = scale_numerical_features(data, [‘longitude’, ‘latitude’, ‘review_count_x’, ‘review_count’])<br>
# Encode categorical features<br>
data = encode_categorical_features(data, [‘city’])<br>
return data</p>
<p>def scale_numerical_features(data, features):<br>
“”“Scale numerical features using StandardScaler”“”<br>
scaled_data = data.copy()<br>
scaled_data[features] = scaler.fit_transform(scaled_data[features])<br>
return scaled_data</p>
<p>def encode_categorical_features(data, features):<br>
“”“Encode categorical features using OneHotEncoder”“”<br>
encoded_data = data.copy()<br>
encoded_data[features] = encoder.fit_transform(encoded_data[features])<br>
return encoded_data</p>
<p>def create_bipartite_graphs(data):<br>
“”“Create user-restaurant and user-category bipartite graphs.”“”<br>
user_restaurant_graph = nx.Graph()<br>
user_category_graph = nx.Graph()</p>
<pre><code># 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
</code></pre>
<p>def add_node_features(user_restaurant_graph, user_category_graph, data):<br>
“”“Add node features to the bipartite graphs.”“”<br>
# Add ‘review_count’ as a node feature<br>
user_review_count = dict(zip(data[‘user_id’], data[‘review_count_x’]))<br>
restaurant_review_count = dict(zip(data[‘business_id’], data[‘review_count’]))<br>
nx.set_node_attributes(user_restaurant_graph, user_review_count, name=‘review_count_user’)<br>
nx.set_node_attributes(user_restaurant_graph, restaurant_review_count, name=‘review_count_restaurant’)</p>
<pre><code># 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')
</code></pre>
<p>def add_edge_features(user_restaurant_graph, user_category_graph, data):<br>
“”“Add edge features to the bipartite graphs.”“”<br>
# Add ‘text’ as an edge feature<br>
user_restaurant_text = {(row[‘user_id’], row[‘business_id’]): row[‘text’] for _, row in data.iterrows()}<br>
user_category_text = {(row[‘user_id’], category): row[‘text’] for _, row in data.iterrows() for category in row[‘categories’]}<br>
nx.set_edge_attributes(user_restaurant_graph, user_restaurant_text, name=‘review_text’)<br>
nx.set_edge_attributes(user_category_graph, user_category_text, name=‘review_text’)</p>
<pre><code># 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)
</code></pre>
<p>def encode_node_features(node_data):<br>
“”“Encode node features using node attributes and review text.”“”<br>
<span class="hashtag-raw">#print</span>(f"Node data: {node_data}")  # Add this line<br>
node_attributes = <span class="chcklst-box fa fa-square-o fa-fw"></span></p>
<pre><code># 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)
</code></pre>
<p>def create_pyg_data(graph, node_type1, node_type2, node_type3=None):<br>
“”“Create a PyG Data or HeteroData object from a bipartite graph.”“”<br>
data = HeteroData()</p>
<pre><code># 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
</code></pre>
<p>def create_combined_data(user_restaurant_graph, user_category_graph, preprocessed_data):<br>
“”“Create the combined HeteroData object.”“”<br>
# Add node features<br>
add_node_features(user_restaurant_graph, user_category_graph, preprocessed_data)</p>
<pre><code># 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
</code></pre>
<p>def main():<br>
# Load the dataset<br>
data = pd.read_csv(“/kaggle/working/subset_review_merged1.csv”)</p>
<pre><code># 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
</code></pre>
<p>if <strong>name</strong> == ‘<strong>main</strong>’:<br>
combined_graph, combined_data = main()</p>
<pre><code># 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)
</code></pre>
<p>This is the combined_data even after the split:<br>
data: HeteroData(<br>
user={ x=[583, 53] },<br>
restaurant={ x=[142, 53] },<br>
category={ x=[122, 53] },<br>
(user, to, restaurant)={<br>
edge_index=[2, 3669],<br>
edge_attr=[3669, 50],<br>
edge_weight=[592]<br>
},<br>
(user, to, category)={<br>
edge_index=[2, 3669],<br>
edge_attr=[3669, 50],<br>
edge_weight=[3077]<br>
},<br>
(restaurant, rev_to, user)={<br>
edge_index=[2, 3669],<br>
edge_attr=[3669, 50]<br>
},<br>
(category, rev_to, user)={<br>
edge_index=[2, 3669],<br>
edge_attr=[3669, 50]<br>
}<br>
)</p>
<p>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?</p>