diff --git a/DRAMSys/traceAnalyzer/presentation/tracedrawingproperties.cpp b/DRAMSys/traceAnalyzer/presentation/tracedrawingproperties.cpp
index 930f0e87..95c0a4cd 100644
--- a/DRAMSys/traceAnalyzer/presentation/tracedrawingproperties.cpp
+++ b/DRAMSys/traceAnalyzer/presentation/tracedrawingproperties.cpp
@@ -103,6 +103,9 @@ void TraceDrawingProperties::updateTracePlotLines(QTreeWidget *treeWidget)
 
 void TraceDrawingProperties::addRankGroup(unsigned int rank)
 {
+    if (rank >= numberOfRanks)
+        return;
+
     std::shared_ptr<TracePlotFirstRankLine> firstRankLine;
 
     for (unsigned int group = groupsPerRank; group--;)
diff --git a/DRAMSys/traceAnalyzer/presentation/traceselector.cpp b/DRAMSys/traceAnalyzer/presentation/traceselector.cpp
index 9ca682fb..434f7ac2 100644
--- a/DRAMSys/traceAnalyzer/presentation/traceselector.cpp
+++ b/DRAMSys/traceAnalyzer/presentation/traceselector.cpp
@@ -124,6 +124,29 @@ void TraceSelectedTree::init(unsigned int groupsPerRank, unsigned int banksPerGr
     this->banksPerGroup = banksPerGroup;
 }
 
+std::vector<QTreeWidgetItem *> TraceSelectedTree::clonedItems() const
+{
+    std::vector<QTreeWidgetItem *> clonedTopLevelItems;
+
+    for (unsigned int i = 0; i < topLevelItemCount(); i++)
+    {
+        clonedTopLevelItems.push_back(topLevelItem(i)->clone());
+    }
+
+    return clonedTopLevelItems;
+}
+
+void TraceSelectedTree::setItems(const std::vector<QTreeWidgetItem *> &items)
+{
+    for (unsigned int i = topLevelItemCount(); i--;)
+        takeTopLevelItem(i);
+
+    for (auto item : items)
+        addTopLevelItem(item);
+
+    emit itemsChanged();
+}
+
 QMimeData *TraceSelectedTree::mimeData(const QList<QTreeWidgetItem *> items) const
 {
     QByteArray traceLineData;
@@ -346,6 +369,16 @@ void TraceSelector::setUpTreeWidget(QTreeWidget *treeWidget)
     );
 }
 
+std::vector<QTreeWidgetItem *> TraceSelector::getTraceSelectorState() const
+{
+    return selectedTree->clonedItems();
+}
+
+void TraceSelector::restoreTraceSelectorState(const std::vector<QTreeWidgetItem *> &items)
+{
+    selectedTree->setItems(items);
+}
+
 TraceSelectorTreeItem *TraceSelectorTreeItem::createRankParentItem(unsigned int rank,unsigned int groupsPerRank,
                                                                   unsigned int banksPerGroup)
 {
@@ -409,3 +442,19 @@ QString TraceSelectorTreeItem::getLabel(unsigned int rank, unsigned int bankgrou
 {
     return "RA" + QString::number(rank) + " BG" + QString::number(bankgroup) + " BA" + QString::number(bank);
 }
+
+QTreeWidgetItem *TraceSelectorTreeItem::clone() const
+{
+    TraceSelectorTreeItem *cloned = new TraceSelectorTreeItem(*this);
+
+    std::vector<TraceSelectorTreeItem *> clonedChildren;
+    for (unsigned int i = 0; i < childCount(); i++)
+    {
+        clonedChildren.push_back(static_cast<TraceSelectorTreeItem *>(child(i)->clone()));
+    }
+
+    for (auto child : clonedChildren)
+        cloned->addChild(child);
+
+    return cloned;
+}
diff --git a/DRAMSys/traceAnalyzer/presentation/traceselector.h b/DRAMSys/traceAnalyzer/presentation/traceselector.h
index 7d120da2..6a2a9fd3 100644
--- a/DRAMSys/traceAnalyzer/presentation/traceselector.h
+++ b/DRAMSys/traceAnalyzer/presentation/traceselector.h
@@ -70,6 +70,9 @@ public:
     explicit TraceSelectedTree(QWidget *parent = nullptr);
     void init(unsigned int groupsPerRank, unsigned int banksPerGroup);
 
+    std::vector<QTreeWidgetItem *> clonedItems() const;
+    void setItems(const std::vector<QTreeWidgetItem *> &items);
+
 Q_SIGNALS:
     void itemsChanged();
 
@@ -108,6 +111,9 @@ public:
     static QString getLabel(unsigned int rank, unsigned int bankgroup, unsigned int bank);
 
     const unsigned int rank = 0, bankgroup = 0, bank = 0;
+
+protected:
+    QTreeWidgetItem *clone() const override;
 };
 
 class TraceSelector : public QWidget
@@ -118,6 +124,9 @@ public:
     explicit TraceSelector(QWidget *parent = nullptr);
     void init(TracePlot *traceplot);
 
+    std::vector<QTreeWidgetItem *> getTraceSelectorState() const;
+    void restoreTraceSelectorState(const std::vector<QTreeWidgetItem *> &items);
+
 signals:
     void selectedTreeChanged(QTreeWidget *treeWidget);
 
diff --git a/DRAMSys/traceAnalyzer/traceanalyzer.cpp b/DRAMSys/traceAnalyzer/traceanalyzer.cpp
index f3142fa1..d771a4ac 100644
--- a/DRAMSys/traceAnalyzer/traceanalyzer.cpp
+++ b/DRAMSys/traceAnalyzer/traceanalyzer.cpp
@@ -178,6 +178,7 @@ void TraceAnalyzer::reloadTab(int index)
     QString traceFile = traceFileTab->getPathToTraceFile();
     traceTime time = traceFileTab->getCurrentTraceTime();
     traceTime zoomLevel = traceFileTab->getZoomLevel();
+    auto treeItems = traceFileTab->saveTraceSelectorState();
 
     if (!traceFileTab->close())
         return;
@@ -188,6 +189,7 @@ void TraceAnalyzer::reloadTab(int index)
     traceFileTab = createTraceFileTab(traceFile);
     traceFileTab->setZoomLevel(zoomLevel);
     traceFileTab->navigateToTime(time);
+    traceFileTab->restoreTraceSelectorState(treeItems);
 
     ui->traceFileTabs->insertTab(index, traceFileTab, QFileInfo(traceFile).baseName());
     ui->traceFileTabs->setCurrentIndex(index);
diff --git a/DRAMSys/traceAnalyzer/tracefiletab.cpp b/DRAMSys/traceAnalyzer/tracefiletab.cpp
index 4438da91..ee568c33 100644
--- a/DRAMSys/traceAnalyzer/tracefiletab.cpp
+++ b/DRAMSys/traceAnalyzer/tracefiletab.cpp
@@ -152,7 +152,7 @@ void TraceFileTab::setUpFileWatcher(QString path)
 
 void TraceFileTab::setUpTraceSelector()
 {
-    TraceSelector *selector = ui->tabCustomizePlot;
+    TraceSelector *selector = static_cast<TraceSelector *>(ui->tabCustomizePlot);
     connect(selector, &TraceSelector::selectedTreeChanged,
             ui->traceplot, &TracePlot::updateTracePlotLines);
     selector->init(ui->traceplot);
@@ -250,6 +250,18 @@ void TraceFileTab::setZoomLevel(traceTime zoomLevel)
     tracescroller->tracePlotZoomChanged();
 }
 
+std::vector<QTreeWidgetItem *> TraceFileTab::saveTraceSelectorState() const
+{
+    TraceSelector *selector = static_cast<TraceSelector *>(ui->tabCustomizePlot);
+    return selector->getTraceSelectorState();
+}
+
+void TraceFileTab::restoreTraceSelectorState(const std::vector<QTreeWidgetItem *> &items)
+{
+    TraceSelector *selector = static_cast<TraceSelector *>(ui->tabCustomizePlot);
+    selector->restoreTraceSelectorState(items);
+}
+
 class ItemDelegate: public QItemDelegate
 {
 public:
diff --git a/DRAMSys/traceAnalyzer/tracefiletab.h b/DRAMSys/traceAnalyzer/tracefiletab.h
index 50703b03..aee1acf0 100644
--- a/DRAMSys/traceAnalyzer/tracefiletab.h
+++ b/DRAMSys/traceAnalyzer/tracefiletab.h
@@ -78,6 +78,9 @@ public:
     traceTime getZoomLevel() const;
     void setZoomLevel(traceTime zoomLevel);
 
+    std::vector<QTreeWidgetItem *> saveTraceSelectorState() const;
+    void restoreTraceSelectorState(const std::vector<QTreeWidgetItem *> &items);
+
 protected:
     void closeEvent(QCloseEvent *event) override;