From a9498fdde9df3a81006cf6c4b8007bac357fb606 Mon Sep 17 00:00:00 2001
From: leschultz <laneenriqueschultz@gmail.com>
Date: Wed, 17 Apr 2024 14:46:25 -0500
Subject: [PATCH] Fixed bug

---
 src/madml/plots.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/madml/plots.py b/src/madml/plots.py
index 877dc63..f53f284 100644
--- a/src/madml/plots.py
+++ b/src/madml/plots.py
@@ -548,13 +548,13 @@ def __init__(
 
         # Domain prediction columns
         pred_cols = [i for i in self.df.columns if 'Domain Prediction' in i]
-        pred_cols = [i.split(' (')[0] for i in pred_cols]
-        self.pred_cols = pred_cols
 
         if df_confusion is None:
             self.df_confusion = df[self.domains+pred_cols+['splitter']]
+            self.pred_cols = pred_cols
         else:
             self.df_confusion = df_confusion
+            self.pred_cols = [i.split(' (')[0] for i in pred_cols]
 
     def generate(self):