get axes

dianna-ai · Dec 15, 2023 · 7106357 · 7106357
1 parent b2e5783
commit 7106357
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/dianna/visualization/tabular.py b/dianna/visualization/tabular.py
@@ -31,17 +31,17 @@ def plot_tabular(
     """
     if not num_features:
         num_features = len(x)
-    fig = plt.figure()
+    fig, ax = plt.subplots()
     abs_values = [abs(i) for i in x]
     top_values = [x for _, x in sorted(zip(abs_values, x), reverse=True)][:num_features]
     top_features = [x for _, x in sorted(zip(abs_values, y), reverse=True)][
         :num_features
     ]
 
     colors = ["r" if x >= 0 else "b" for x in top_values]
-    plt.barh(top_features, top_values, color=colors)
-    plt.xlabel(x_label)
-    plt.ylabel(y_label)
+    ax.barh(top_features, top_values, color=colors)
+    ax.set_xlabel(x_label)
+    ax.set_ylabel(y_label)
 
     if show_plot:
         plt.show()