From 8a430018d75c8c6ee54b4c21674db9239679f739 Mon Sep 17 00:00:00 2001
From: Nick Touran <ntouran@terrapower.com>
Date: Sat, 26 Oct 2024 17:59:26 -0400
Subject: [PATCH] Avoiding closing plots that are meant to be interactive
 (#1978)

The plt.show() options are intended for:

1. Interactive use on ipython or notebooks
2. Automated use during sphinx-gallery generation

Running plt.show() followed by plt.close() prevents at
least the 2nd use case from working, as seen in #1977.

This change puts the close() only in the code branches
that are expected to be run non-interactively, e.g.
compute code. That allows the system to keep memory
clear during big runs while still allowing interactive
and gallery usage.

Fixes #1977
---
 armi/nuclearDataIO/xsNuclides.py                |  6 ++----
 armi/reactor/converters/blockConverters.py      |  2 +-
 armi/utils/plotting.py                          | 16 +++++++++-------
 doc/gallery-src/analysis/run_hexReactorToRZ.py  |  1 -
 doc/gallery-src/framework/run_fuelManagement.py |  2 --
 doc/gallery-src/framework/run_reactorFacemap.py |  1 -
 6 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/armi/nuclearDataIO/xsNuclides.py b/armi/nuclearDataIO/xsNuclides.py
index e11296a7e..c2c976af4 100644
--- a/armi/nuclearDataIO/xsNuclides.py
+++ b/armi/nuclearDataIO/xsNuclides.py
@@ -240,11 +240,10 @@ def plotScatterMatrix(scatterMatrix, scatterTypeLabel="", fName=None):
     pyplot.colorbar()
     if fName:
         pyplot.savefig(fName)
+        pyplot.close()
     else:
         pyplot.show()
 
-    pyplot.close()
-
 
 def plotCompareScatterMatrix(scatterMatrix1, scatterMatrix2, fName=None):
     """Compares scatter matrices graphically between libraries."""
@@ -260,7 +259,6 @@ def plotCompareScatterMatrix(scatterMatrix1, scatterMatrix2, fName=None):
     pyplot.colorbar()
     if fName:
         pyplot.savefig(fName)
+        pyplot.close()
     else:
         pyplot.show()
-
-    pyplot.close()
diff --git a/armi/reactor/converters/blockConverters.py b/armi/reactor/converters/blockConverters.py
index dac9a9588..fc9787151 100644
--- a/armi/reactor/converters/blockConverters.py
+++ b/armi/reactor/converters/blockConverters.py
@@ -488,9 +488,9 @@ def plotConvertedBlock(self, fName=None):
         fig.tight_layout()
         if fName:
             plt.savefig(fName)
+            plt.close()
         else:
             plt.show()
-        plt.close()
         return fName
 
 
diff --git a/armi/utils/plotting.py b/armi/utils/plotting.py
index 8e713eeb1..0726c5d9b 100644
--- a/armi/utils/plotting.py
+++ b/armi/utils/plotting.py
@@ -178,10 +178,10 @@ def update(i):
 
     if fName:
         plt.savefig(fName, dpi=150)
+        plt.close()
     else:
         plt.show()
 
-    plt.close()
     return fName
 
 
@@ -402,13 +402,16 @@ def plotFaceMap(
                 "Cannot update facemap at {0}: IOError. Is the file open?"
                 "".format(fName)
             )
+        plt.close(fig)
     elif referencesToKeep:
         # Don't show yet, since it will be updated.
         return fName
     else:
+        # Never close figures after a .show()
+        # because they're being used interactively e.g.
+        # in a live tutorial or by the doc gallery
         plt.show()
 
-    plt.close(fig)
     return fName
 
 
@@ -823,8 +826,8 @@ def plotAssemblyTypes(
     if fileName:
         fig.savefig(fileName)
         runLog.debug("Writing assem layout {} in {}".format(fileName, os.getcwd()))
+        plt.close(fig)
 
-    plt.close(fig)
     return fig
 
 
@@ -1092,11 +1095,11 @@ def getTable(self):
             os.path.abspath(fName),
             report.FLUX_PLOT,
         )
+        plt.close()
     else:
+        # Never close interactive plots
         plt.show()
 
-    plt.close()
-
 
 def makeHistogram(x, y):
     """
@@ -1503,7 +1506,6 @@ def plotNucXs(
 
     if fName:
         plt.savefig(fName)
+        plt.close()
     elif not noShow:
         plt.show()
-
-    plt.close()
diff --git a/doc/gallery-src/analysis/run_hexReactorToRZ.py b/doc/gallery-src/analysis/run_hexReactorToRZ.py
index 89b3e6279..1baa4c7e2 100644
--- a/doc/gallery-src/analysis/run_hexReactorToRZ.py
+++ b/doc/gallery-src/analysis/run_hexReactorToRZ.py
@@ -55,4 +55,3 @@
 figs = converter.plotConvertedReactor()
 
 plt.show()
-plt.close()
diff --git a/doc/gallery-src/framework/run_fuelManagement.py b/doc/gallery-src/framework/run_fuelManagement.py
index 44a18df1c..d355bbaf0 100644
--- a/doc/gallery-src/framework/run_fuelManagement.py
+++ b/doc/gallery-src/framework/run_fuelManagement.py
@@ -30,7 +30,6 @@
 # sphinx_gallery_thumbnail_number = 2
 import math
 
-import matplotlib.pyplot as plt
 
 from armi import configure
 from armi.physics.fuelCycle import fuelHandlers
@@ -73,4 +72,3 @@
 
 # show final burnup distribution
 plotting.plotFaceMap(reactor.core, param="percentBu")
-plt.close()
diff --git a/doc/gallery-src/framework/run_reactorFacemap.py b/doc/gallery-src/framework/run_reactorFacemap.py
index 7b961d461..457aa3f5f 100644
--- a/doc/gallery-src/framework/run_reactorFacemap.py
+++ b/doc/gallery-src/framework/run_reactorFacemap.py
@@ -33,4 +33,3 @@
     b.p.pdens = x**2 + y**2 + z**2
 
 plotting.plotFaceMap(reactor.core, param="pdens", labelFmt="{0:.1e}")
-plotting.close()