Two-Key plot color map fixes and tests

Author: Zan Vrabic

niaarm/visualize.py

@@ -493,7 +493,7 @@ def prepare_data(rules, metrics):
         return data_frame
 
     # Check if one or more rules
-    if not hasattr(rules, "data"):
+    if not hasattr(rules, "data") and not isinstance(rules, list):
         rules = [rules]
 
     # Prepare the data
@@ -528,16 +528,21 @@ def prepare_data(rules, metrics):
 
         # Map each order to a unique color
         unique_orders = sorted(df["order"].unique())
-        color_map = plt.cm.get_cmap("tab10", len(unique_orders))  
-        color_mapping = {order: color_map(i) for i, order in enumerate(unique_orders)}
+        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(
-                subset[metrics[0]],
-                subset[metrics[1]],
-                label=f"Order {order}",
+                x_data,
+                y_data,
+                label=order,
                 color=color_mapping[order],
                 alpha=0.7
             )

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.")