Merge pull request #148 from vrabiczan/TwoKey_plot

Two key plot

Author: firefly-cpp

examples/visualization_examples/two_key_plot/data_developer_salary_two_key_plot.py

@@ -0,0 +1,39 @@
+from examples.visualization_examples.prepare_datasets import get_data_developer_salary_data
+from niaarm import Dataset, get_rules
+from niaarm.visualize import two_key_plot
+
+# Get prepared data developer salary data
+arm_df = get_data_developer_salary_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
+)
+
+# 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 = two_key_plot(rules=rules, metrics=metrics, interactive=True)
+fig.show()

niaarm/visualize.py

@@ -448,3 +448,110 @@ def create_plot_data(data_frame):
         plt.grid(which="both", color="grey", linestyle="-", linewidth=0.5)
 
         return plt
+
+def two_key_plot(rules, metrics, interactive=False):
+    """
+    Visualize rules as a two key plot with two primary metrics (support, confidence) and rule order.
+    
+    Args:
+        rules (Rule): Association rule or rules to visualize.
+        metrics (tuple): Two metrics to display on the x and y axes. 'order' will be used for point color.
+        interactive (bool): Make plot interactive. Default: False.
+    
+    Returns:
+        Figure or plot.
+    """
+
+    # Ensure exactly two metrics for the axes
+    if len(metrics) != 2:
+        raise ValueError("Please provide exactly two metrics for a two-key plot.")
+
+    # Function to prepare the data
+    def prepare_data(rules, metrics):
+        data = {
+            "rule": [],
+            metrics[0]: [],
+            metrics[1]: [],
+            "order": []  # Store rule order (length)
+        }
+
+        for rule in rules:
+            data["rule"].append(rule.__repr__())
+            data[metrics[0]].append(getattr(rule, metrics[0]))
+            data[metrics[1]].append(getattr(rule, metrics[1]))
+            
+            # Calculate order dynamically as the total number of items in antecedent and consequent
+            if hasattr(rule, 'antecedent') and hasattr(rule, 'consequent'):
+                rule_order = len(rule.antecedent) + len(rule.consequent)
+            else:
+                rule_order = 0  # Fallback if structure is missing
+            
+            data["order"].append(rule_order)
+
+        # Return as DataFrame
+        data_frame = pd.DataFrame(data)
+        return data_frame
+
+    # Check if one or more rules
+    if not hasattr(rules, "data") and not isinstance(rules, list):
+        rules = [rules]
+
+    # Prepare the data
+    df = prepare_data(rules, metrics)
+
+    # Interactive plot using Plotly
+    if interactive:
+        title = f'Interactive two-key plot for {len(rules)} rules' \
+            if len(rules) > 1 else "Interactive two-key plot for rule"
+        
+        # Create figure
+        fig = px.scatter(
+            data_frame=df,
+            x=metrics[0],
+            y=metrics[1],
+            color=df["order"].astype(str),
+            hover_name="rule",
+            title=title,
+            labels={"color": "order"},
+            color_discrete_sequence=px.colors.qualitative.Plotly
+        )
+        fig.update_layout(
+            xaxis_title=metrics[0],
+            yaxis_title=metrics[1],
+            legend_title = "Order"
+        )
+        return fig
+
+    # Static plot using Matplotlib
+    else:
+        plt.figure(figsize=(12, 8))
+
+        # Map each order to a unique color
+        unique_orders = sorted(df["order"].unique())
+        color_map = plt.colormaps.get_cmap("Set1")
+        color_indices = np.linspace(0, 1, len(unique_orders)) 
+        colors = [color_map(i) for i in color_indices]
+        color_mapping = {order: colors[i] for i, order in enumerate(unique_orders)}
+
+        # Plot each order separately for discrete colors
+        for order in unique_orders:
+            subset = df[df["order"] == order]
+            x_data = np.array(subset[metrics[0]].tolist())
+            y_data = np.array(subset[metrics[1]].tolist())
+			
+            plt.scatter(
+                x_data,
+                y_data,
+                label=order,
+                color=color_mapping[order],
+                alpha=0.7
+            )
+
+        # Add legend and labels
+        plt.title(f'Two-key plot for {len(rules)} rules')
+        plt.xlabel(metrics[0])
+        plt.ylabel(metrics[1])
+        plt.legend(title="Order")
+        plt.grid(True)
+        return plt
+        

tests/test_two_key_plot.py

@@ -0,0 +1,46 @@
+from unittest import TestCase
+import matplotlib.pyplot as plt
+import pandas as pd
+from niaarm.visualize import two_key_plot
+
+class Rule:
+    def __init__(self, antecedent, consequent, support, confidence):
+        self.antecedent = antecedent
+        self.consequent = consequent
+        self.support = support
+        self.confidence = confidence
+
+    def __repr__(self):
+        return f"Rule({self.antecedent} -> {self.consequent})"
+
+class TestTwoKeyPlot(TestCase):
+    
+    @classmethod
+    def setUpClass(cls):
+        cls.rule1 = Rule(antecedent=["A", "B"], consequent=["C"], support=0.3, confidence=0.8)
+        cls.rule2 = Rule(antecedent=["D"], consequent=["E", "F"], support=0.5, confidence=0.7)
+        cls.rule3 = Rule(antecedent=["G", "H"], consequent=["I"], support=0.2, confidence=0.9)
+
+        cls.rules = [cls.rule1, cls.rule2, cls.rule3]  # Ensure rules are available to all tests
+
+    def test_two_key_plot(self):
+        metrics = ("support", "confidence")
+
+        plot = two_key_plot(self.rules, metrics, interactive=False)
+
+        # Verify that the return type is Matplotlib's pyplot
+        self.assertIs(plot, plt)
+
+        # Ensure a figure is created
+        self.assertTrue(plt.gcf().axes, "No axes found in the generated plot.")
+
+    def test_invalid_metrics(self):
+        with self.assertRaises(ValueError):
+            two_key_plot(self.rules, ("support",), interactive=False)
+
+    def test_interactive_plot(self):
+        metrics = ("support", "confidence")
+        fig = two_key_plot(self.rules, metrics, interactive=True)
+
+        # Verify that a Plotly figure is returned
+        self.assertEqual(fig.__class__.__name__, "Figure", "Expected a Plotly figure but got a different type.")