Sankey diagram v1 and example

Author: Zan Vrabic

examples/visualization_examples/sankey_diagram/weather_data_sankey_diagram.py

@@ -0,0 +1,45 @@
+from examples.visualization_examples.prepare_datasets import get_weather_data
+from niaarm import Dataset, get_rules
+from niaarm.visualize import sankey_diagram
+
+# Get prepared weather data
+arm_df = get_weather_data()
+
+# Prepare Dataset
+dataset = Dataset(
+    path_or_df=arm_df,
+    delimiter=","
+)
+
+# Get rules
+metrics = ("support", "confidence")
+rules, run_time = get_rules(
+    dataset=dataset,
+    algorithm="DifferentialEvolution",
+    metrics=metrics,
+    max_evals=500
+)
+
+# Add lift after the rules have been generated
+# Cannot be in metrics before because get_rules metrics doesn't contain lift, therefore we need to add after
+metrics = list(metrics)
+metrics.append("lift")
+metrics = tuple(metrics)
+
+# Sort rules
+rules.sort(by="support")
+# Print rule information
+print("\nRules:")
+print(rules)
+print(f'\nTime to generate rules: {f"{run_time:.3f}"} seconds')
+print("\nRule information: ", rules[3])
+print("Antecedent: ", rules[3].antecedent)
+print("Consequent: ", rules[3].consequent)
+print("Confidence: ", rules[3].confidence)
+print("Support: ", rules[3].support)
+print("Lift: ", rules[3].lift)
+print("\nMetrics:", metrics)
+
+# Visualize scatter plot
+fig = sankey_diagram(rules=rules, interestingness_measure="support", M=4)
+fig.show()

niaarm/visualize.py

@@ -3,8 +3,10 @@
 from matplotlib.colors import Normalize
 import numpy as np
 import plotly.express as px
+import plotly.graph_objects as go
 import pandas as pd
 from sklearn.cluster import KMeans
+from itertools import combinations
 
 
 def hill_slopes(rule, transactions):
@@ -554,4 +556,115 @@ def prepare_data(rules, metrics):
         plt.legend(title="Order")
         plt.grid(True)
         return plt
-        
+
+
+def sankey_diagram(rules, interestingness_measure, M=4):
+    """
+    Visualize rules as a sankey diagram.
+    
+    Args:
+        rules (Rule): Association rule or rules to visualize.
+        interestingness_measures (str): Interestingness measure Z = {supp, cons, lift},reflecting the quality of a particular connection.
+        m (int): Maximum number of rules to be selected for visualization. Default: 4
+    
+    Returns:
+        Figure or plot.
+    """
+
+    
+    def compute_similarity(rule1, rule2):
+        """Compute similarity between two rules."""
+        ant_inter = len(set(str(rule1.antecedent)) & set(str(rule2.antecedent)))
+        ant_union = len(set(str(rule1.antecedent)) | set(str(rule2.antecedent)))
+        con_inter = len(set(str(rule1.consequent)) & set(str(rule2.consequent)))
+        con_union = len(set(str(rule1.consequent)) | set(str(rule2.consequent)))
+        return (ant_inter + con_inter) / (ant_union + con_union)
+
+    def build_adjacency_matrix(rules):
+        size = len(rules)
+        adjacency_matrix = np.zeros((size, size))
+
+        for i, j in combinations(range(size), 2):
+            similarity = compute_similarity(rules[i], rules[j])
+            adjacency_matrix[i, j] = similarity
+            adjacency_matrix[j, i] = similarity
+
+        return adjacency_matrix
+    
+    def knapsack_selection(adj_matrix, rules, M):
+        fitness_scores = np.array([rule.fitness for rule in rules])
+        N = len(rules)
+        weights = np.ones(N)
+        similarity_weight = 1.0
+        fitness_weight = 0.5
+        combined_profits = similarity_weight * np.sum(adj_matrix) + fitness_weight * fitness_scores
+    
+        selected = np.zeros(N, dtype=int)
+    
+        # Initialize DP table
+        dp = np.zeros((N + 1, M + 1))
+        for i in range(1, N + 1):
+            for w in range(1, M + 1):
+                if weights[i - 1] <= w:
+                    dp[i, w] = max(dp[i - 1, w], dp[i - 1, w - 1] + combined_profits[i - 1])
+                else:
+                    dp[i, w] = dp[i - 1, w]
+    
+        # Backtrack to find selected rules
+        w = M
+        for i in range(N, 0, -1):
+            if dp[i, w] != dp[i - 1, w]:
+                selected[i - 1] = 1
+                w -= 1
+    
+        selected_rules = [rules[i] for i in range(N) if selected[i]]
+
+        return selected_rules
+
+    def prepare_data(rules, M, interestingness_measure):
+        adj_matrix = build_adjacency_matrix(rules)
+        selected_rules = knapsack_selection(adj_matrix, rules, M)
+
+        sources=[]
+        targets=[] 
+        values=[] 
+        labels=[]
+        node_indices = {}
+
+        for rule in selected_rules:
+            for antecedent in rule.antecedent:
+                if str(antecedent) not in node_indices:
+                    node_indices[str(antecedent)] = len(labels)
+                    labels.append(str(antecedent))
+                sources.append(node_indices[str(antecedent)])
+            
+            for consequent in rule.consequent:
+                if str(consequent) not in node_indices:
+                    node_indices[str(consequent)] = len(labels)
+                    labels.append(str(consequent))
+                targets.append(node_indices[str(consequent)])
+            
+            measure_value = getattr(rule, interestingness_measure, rule.support) #default support
+            values.append(measure_value) 
+
+        return labels, sources, targets, values
+
+    labels, sources, targets, values = prepare_data(rules, M, interestingness_measure)
+
+    fig = go.Figure(go.Sankey(
+        node=dict(
+            pad=15, 
+            thickness=20, 
+            line=dict(color='black', width=0.5),
+            label=labels
+        ),
+        link=dict(
+            source=sources,
+            target=targets,
+            value=values
+        )
+    ))
+    fig.update_layout(title_text=f'Sankey Diagram of Association Rules ({interestingness_measure})', font_size=10)
+    
+    return fig       
+