Visualization of Heterograph Data

[kyuwanchoi]

Hi DGL community members,

I created a hetero_graph data using the following code and would like to visualize it. Is there any way to visualize hetero_graph data?

import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import seaborn as sns
import dgl
import torch
import dgl.nn as dglnn
import torch.nn as nn
import torch.nn.functional as F
import seaborn as sns
from matplotlib import rcParams

path = 's3://gilead-gna-corp-ops-dev-us-west-2-gna-datascience/credit_card_fraud_detection/'
df1 = pd.read_csv('/kaggle/input/fraud-detection/fraudTrain.csv')
df2 = pd.read_csv('/kaggle/input/fraud-detection/fraudTest.csv')

df1['train_mask'] = 1
df1['test_mask'] = 0
df2['train_mask'] = 0
df2['test_mask'] = 1
df = pd.concat([df1, df2])

def encode_data(input):

    customer_node = set(df['cc_num'])
    merchant_node = set(df['merchant'])
    transaction_node = set(df['trans_num'])
    category = set(df['category'])
    city = set(df['city'])

    mapping_c = dict(zip(customer_node, range(len(customer_node))))
    mapping_m = dict(zip(merchant_node, range(len(merchant_node))))                 
    mapping_t = dict(zip(transaction_node, range(len(df))))
    mapping_cat = dict(zip(category, range(len(df))))
    mapping_city = dict(zip(city, range(len(df))))
    
    output = input.copy()
    output['customer_id'] = output.cc_num.apply(lambda x: mapping_c[x])
    output['merchant_id'] = output.merchant.apply(lambda x: mapping_m[x])
    output['transaction_id'] = output.trans_num.apply(lambda x: mapping_t[x])
    output['category_id'] = output.category.apply(lambda x: mapping_cat[x])
    output['city_id'] = output.city.apply(lambda x: mapping_city[x])

    return output

encoded_df = encode_data(df)

fraud_item = encoded_df[encoded_df['is_fraud'] == 1]
non_fraud_item = encoded_df[encoded_df['is_fraud'] == 0]

fraud_buy_edge = [tuple(x) for x in fraud_item[['customer_id', 'merchant_id']].to_numpy()]
non_buy_edge = [tuple(x) for x in non_fraud_item[['customer_id', 'merchant_id']].to_numpy()]
fraud_sell_edge = [tuple(x) for x in fraud_item[['merchant_id', 'customer_id']].to_numpy()]
non_sell_edge = [tuple(x) for x in non_fraud_item[['merchant_id', 'customer_id']].to_numpy()]

hetero_graph = dgl.heterograph({('client', 'is_fraud_buy', 'merchant'):fraud_buy_edge,
                     ('client', 'not_fraud_buy', 'merchant'): non_buy_edge,
                     ('merchant', 'is_fraud_sell', 'client'): fraud_sell_edge,
                     ('merchant', 'not_fraud_sell', 'client'): non_sell_edge})

client_node = encoded_df.copy()
client_amt = client_node.groupby(['customer_id']).sum().amt
client_node = client_node.groupby(['customer_id']).max()
client_node['amt'] = client_amt

merchant_node = encoded_df.copy()
merchant_amt = merchant_node.groupby(['merchant_id']).sum().amt
merchant_node = merchant_node.groupby(['merchant_id']).max()
merchant_node['amt'] = merchant_amt

hetero_graph.nodes['client'].data['feature'] = torch.tensor(client_node[['is_fraud','amt']].to_numpy()).float() #default: is_fraud was included
hetero_graph.nodes['merchant'].data['feature'] = torch.tensor(merchant_node[['is_fraud','amt']].to_numpy()).float()
                                       
hetero_graph.edges['is_fraud_buy'].data['train_mask'] = torch.Tensor(fraud_item['train_mask'].to_numpy()).bool()
hetero_graph.edges['is_fraud_buy'].data['test_mask'] = torch.Tensor(fraud_item['test_mask'].to_numpy()).bool()
hetero_graph.edges['not_fraud_buy'].data['train_mask'] = torch.Tensor(non_fraud_item['train_mask'].to_numpy()).bool()
hetero_graph.edges['not_fraud_buy'].data['test_mask'] = torch.Tensor(non_fraud_item['test_mask'].to_numpy()).bool()
hetero_graph.edges['is_fraud_sell'].data['train_mask'] = torch.Tensor(fraud_item['train_mask'].to_numpy()).bool()
hetero_graph.edges['is_fraud_sell'].data['test_mask'] = torch.Tensor(fraud_item['test_mask'].to_numpy()).bool()
hetero_graph.edges['not_fraud_sell'].data['train_mask'] = torch.Tensor(non_fraud_item['train_mask'].to_numpy()).bool()
hetero_graph.edges['not_fraud_sell'].data['test_mask'] = torch.Tensor(non_fraud_item['test_mask'].to_numpy()).bool()

[frozenbugs]

We don’t have dedicated solution for visualize hetero graph, here is solution powered by Gemini, can you try similar ones:

import dgl
import networkx as nx
import matplotlib.pyplot as plt

# Create a sample heterograph (replace this with your actual data)
g = dgl.heterograph({
    ('author', 'writes', 'paper'): ([0, 1, 2], [1, 2, 1]),
    ('author', 'reviews', 'paper'): ([0, 0, 2], [2, 1, 0])
})

# Convert heterograph to NetworkX graph
G = dgl.to_networkx(g)

# Assign different colors for different node types
node_colors = {n: 'blue' if G.nodes[n]['ntype'] == 'author' else 'red' for n in G.nodes}

# Plot the graph with customized node colors
plt.figure(figsize=(8, 6))
nx.draw(G, node_color=node_colors, with_labels=True, font_weight='bold')
plt.show()

[kyuwanchoi]

@fronzenbugs, Thank you for your support. It is very helpful for me.

Best,
Kyu