Separate plots

tests/test.py

@@ -1,4 +1,5 @@
 import sys, logging
+from pathlib import Path
 import unittest
 import warnings
 
@@ -21,7 +22,6 @@
 import tempfile
 from tempfile import TemporaryDirectory
 import os, errno, shutil, glob
-import json
 from os import listdir
 from os.path import isfile, join
 
@@ -5117,6 +5117,7 @@ def test_torch_availability(self):
             self.assertTrue(RMT_Util._svd_vals_accurate is RMT_Util._svd_vals_fast)
 
 
+
     def test_torch_linalg(self):
         # Note that if torch is not available then this will test scipy instead.
         W = np.random.random((50,50))
@@ -5131,6 +5132,21 @@ def test_torch_linalg(self):
         err = np.sum(np.abs(W - W_reconstruct))
         self.assertLess(err, 0.05, f"torch svd absolute reconstruction error was {err}")
 
+	def test_torch_linalg_eig(self):
+		# Note that if torch is not available then this will test scipy instead.
+		W = np.random.random((50,50))
+		W = np.matmul(W, W.T) / 2500
+		L, V = RMT_Util._eig_full_fast(W)
+		W_reconstruct = np.matmul(V.astype("float32"), np.matmul(np.diag(L), np.linalg.inv(V.astype("float32"))))
+		err = np.sum(np.abs(W - W_reconstruct))
+		self.assertLess(err, 0.005, f"torch eig absolute reconstruction error was {err}")
+
+	def test_torch_linalg_svd(self):
+		W = np.random.random((50,100))
+		U, S, Vh = RMT_Util._svd_full_fast(W)
+		W_reconstruct = np.matmul(U.astype("float32"), np.matmul(np.diag(S), Vh[:50,:].astype("float32")))
+		err = np.sum(np.abs(W - W_reconstruct))
+		self.assertLess(err, 0.75, f"torch svd absolute reconstruction error was {err}")
         S_vals_only = RMT_Util._svd_vals_accurate(W)
         err = np.sum(np.abs(S - S_vals_only))
         self.assertLess(err, 0.0005, msg=f"torch svd and svd_vals differed by {err}")
@@ -5252,11 +5268,86 @@ def setUp(self):
         self.model = models.resnet18(weights='ResNet18_Weights.IMAGENET1K_V1')
         self.watcher = ww.WeightWatcher(model=self.model, log_level=logging.WARNING)
 
-    def testPlots(self):
-        """ Simply tests that the plot functions will not generate an exception.
-            Does not guarantee correctness, yet.
-        """
-        self.watcher.analyze(layers=[67], plot=True, randomize=True)
+		self.plotDir = Path("./testPlots")
+
+	def tearDown_plots(self):
+		if not self.plotDir.exists(): return
+
+		for f in self.plotDir.iterdir():
+			f.unlink()
+		self.plotDir.rmdir()
+
+
+	def check_expected_plots(self, plot_figs):
+		if len(plot_figs) == 0:
+			self.assertFalse(self.plotDir.exists())
+			return
+
+		self.assertTrue(self.plotDir.exists())
+
+		figs = list(self.plotDir.iterdir())
+		self.assertEqual(len(figs), len(plot_figs), f"plot={plot_figs} produced {len(figs)} images")
+		self.tearDown_plots()
+
+
+	def testPlots(self):
+		""" Simply tests that the plot functions will not generate an exception.
+			Does not guarantee correctness, yet.
+		"""
+		self.tearDown_plots()	# Sometimes tearDown_plots() doesn't get called when a test fails previously.
+
+		self.watcher.analyze(layers=[67], plot=True, randomize=True, savefig=str(self.plotDir))
+
+		self.assertTrue(self.plotDir.exists(), f"Savefig dir {self.plotDir} should exist after analyze() with plot=True")
+
+		expected_plots = WW_FIT_PL_PLOTS + WW_RANDESD_PLOTS + [WW_PLOT_MPDENSITY]
+		self.check_expected_plots(expected_plots)
+
+
+	def testPlotSingletons(self):
+		self.tearDown_plots()	# Sometimes tearDown_plots() doesn't get called when a test fails previously.
+
+		self.watcher.analyze(layers=[67], plot=[WW_PLOT_DETX], detX=True, savefig=str(self.plotDir))
+		self.check_expected_plots([WW_PLOT_DETX])
+
+		# MPFIT needs Q=1			 \/
+		self.watcher.analyze(layers=[13], plot=[WW_PLOT_MPFIT], mp_fit=True, savefig=str(self.plotDir))
+		self.check_expected_plots([WW_PLOT_MPFIT])
+
+		self.watcher.analyze(layers=[67], plot=[WW_PLOT_MPFIT], mp_fit=True, savefig=str(self.plotDir))
+		self.check_expected_plots([])
+
+		self.watcher.analyze(layers=[67], plot=[WW_PLOT_MPDENSITY], mp_fit=True, savefig=str(self.plotDir))
+		self.check_expected_plots([WW_PLOT_MPDENSITY])
+
+		for plot_fig in WW_FIT_PL_PLOTS:
+			self.watcher.analyze(layers=[67], plot=[plot_fig], savefig=str(self.plotDir))
+			self.check_expected_plots([plot_fig])
+
+		for plot_fig in WW_RANDESD_PLOTS:
+			self.watcher.analyze(layers=[67], plot=[plot_fig], randomize=True, savefig=str(self.plotDir))
+			self.check_expected_plots([plot_fig])
+
+		# Commented for future in case support for this is re-enabled.
+		# for plot_fig in WW_DELTAES_PLOTS:
+		# 	self.watcher.analyze(layers=[67], plot=[plot_fig], deltas=True, savefig=str(self.plotDir))
+		# 	self.check_expected_plots([plot_fig])
+
+
+	def testPlotCombos(self):
+		self.tearDown_plots()	# Sometimes tearDown_plots() doesn't get called when a test fails previously.
+
+		self.watcher.analyze(layers=[67], plot=WW_FIT_PL_PLOTS, savefig=str(self.plotDir))
+		self.check_expected_plots(WW_FIT_PL_PLOTS)
+
+		self.watcher.analyze(layers=[67], plot=WW_RANDESD_PLOTS, randomize=True, savefig=str(self.plotDir))
+		self.check_expected_plots(WW_RANDESD_PLOTS)
+
+		# Commented for future in case support for this is re-enabled.
+		# self.watcher.analyze(layers=[67], plot=WW_DELTAES_PLOTS, deltas=True, savefig=str(self.plotDir))
+		# self.check_expected_plots(WW_DELTAES_PLOTS)
+
+
 
 class Test_Pandas(Test_Base):
     def setUp(self):

weightwatcher/RMT_Util.py

@@ -452,6 +452,9 @@ def plot_density_and_fit(eigenvalues=None, model=None, layer_name="", layer_id=0
 
     # if eigenvalues is None:
     #    eigenvalues = get_eigenvalues(model, weightfile, layer)
+
+    if plot is True:  plot = [WW_PLOT_MPDENSITY]
+    if plot is False: plot = []
 
     if Q == 1:
         to_fit = np.sqrt(eigenvalues)
@@ -463,7 +466,7 @@ def plot_density_and_fit(eigenvalues=None, model=None, layer_name="", layer_id=0
         label = r'$\rho_{emp}(\lambda)$'
         title = " W{} ESD, MP Sigma={:0.3}f" 
 
-    if plot:
+    if WW_PLOT_MPDENSITY in plot:
         plt.hist(to_fit, bins=100, alpha=alpha, color=color, density=True, label=label);
         plt.legend()
 
@@ -491,7 +494,7 @@ def plot_density_and_fit(eigenvalues=None, model=None, layer_name="", layer_id=0
     else:
         x, mp = marchenko_pastur_pdf(x_min, x_max, Q, sigma)
 
-    if plot:
+    if WW_PLOT_MPDENSITY in plot:
         plt.title(title.format(layer_name, sigma))
         plt.plot(x, mp, linewidth=1, color='r', label="MP fit")
 

weightwatcher/constants.py

@@ -154,6 +154,48 @@
 PLOT = 'plot'
 STACKED = 'stacked'
 
+# constants used to indicate which plots should be generated
+WW_PLOT_DETX            = 'detX'
+
+WW_PLOT_VECTOR_METRICS  = 'vectors'
+WW_PLOT_VECTOR_HIST     = 'vector_hist'
+
+WW_PLOT_LOG_DELTAES     = 'log_deltaEs'
+WW_PLOT_DELTAES_LEVELS  = 'deltaEs_levels'
+
+WW_PLOT_MPFIT           = 'mpfit'
+WW_PLOT_MPDENSITY       = 'mpdensity'
+
+WW_PLOT_LOGLOG_ESD      = 'loglog_esd'
+WW_PLOT_LINLIN_ESD      = 'linlin_esd'
+WW_PLOT_LOGLIN_ESD      = 'loglin_esd'
+WW_PLOT_DKS             = 'DKS'
+WW_PLOT_XMIN_ALPHA      = 'xmin_alpha'
+
+WW_PLOT_RANDESD         = 'rand_esd'
+WW_PLOT_LOG_RANDESD     = 'log_rand_esd'
+
+WW_ALL_PLOTS = [
+    WW_PLOT_DETX,
+    WW_PLOT_VECTOR_METRICS, WW_PLOT_VECTOR_HIST,
+    WW_PLOT_LOG_DELTAES, WW_PLOT_DELTAES_LEVELS,
+    WW_PLOT_MPFIT, WW_PLOT_MPDENSITY,
+    WW_PLOT_LOGLOG_ESD, WW_PLOT_LINLIN_ESD, WW_PLOT_LOGLIN_ESD, WW_PLOT_DKS, WW_PLOT_XMIN_ALPHA,
+    WW_PLOT_RANDESD, WW_PLOT_LOG_RANDESD,
+]
+
+WW_FIT_PL_PLOTS = [
+    WW_PLOT_LOGLOG_ESD, WW_PLOT_LINLIN_ESD, WW_PLOT_LOGLIN_ESD, WW_PLOT_DKS, WW_PLOT_XMIN_ALPHA,
+]
+
+WW_RANDESD_PLOTS = [
+    WW_PLOT_RANDESD, WW_PLOT_LOG_RANDESD,
+]
+
+WW_DELTAES_PLOTS = [
+    WW_PLOT_LOG_DELTAES, WW_PLOT_DELTAES_LEVELS,
+]
+
 CHANNELS_STR = 'channels'
 FIRST = 'first'
 LAST = 'last'