study.py

# coding: utf-8
import tools
if __name__ == "__main__":
	raise tools.UserException(f"Module {__file__!r} is not to be used as the main module")

import atexit
import matplotlib.pyplot as plt
import numpy
import pathlib
import pandas
import threading
import types

# Change common font for the default LaTeX one
plt.rcParams["font.family"] = "Latin Modern Roman"
plt.rcParams["font.size"] = 16

# Enable automatic layout adjustments
plt.rcParams["figure.autolayout"] = True

# ---------------------------------------------------------------------------- #
# Common GTK main loop

try:
	import gi
	gi.require_version("Gtk", "3.0")
	from gi.repository import Gtk, Gdk, GLib

	gtk_lazy_lock = threading.Lock()
	gtk_lazy_main = None

	def gtk_run(closure):
		""" Run a closure in the GTK main loop, lazy start it.
		Args:
			closure Closure to run in the main loop
		"""
		global gtk_lazy_lock
		global gtk_lazy_main
		# GTK's main event loop
		def gtk_main():
			# Main loop
			atexit.register(Gtk.main_quit)
			Gtk.main()
		# Lazy-start the loop if necessary
		with gtk_lazy_lock:
			if gtk_lazy_main is None:
				thread = threading.Thread(target=gtk_main, name="gtk_main", daemon=True)
				thread.start()
				gtk_lazy_main = thread
		# Submit the job to the main loop
		GLib.idle_add(closure)

	class _DataFrameDisplayWindow(Gtk.Window):
		""" Display the given data frame in a window.
		"""

		@staticmethod
		def to_string(x):
			""" Convert data to string, special treatment for floats.
			Args:
				x Input data
			Returns:
				Converted data to string
			"""
			if type(x) is float:
				return f"{x:e}"
			return str(x).strip()

		def __init__(self, data, title="Display data"):
			""" Initialize the display window.
			Args:
				data	Data to display
				title Title to use
			"""
			super().__init__(title=title)
			# Make and fill list store
			store = Gtk.ListStore(*([str] * (len(data.columns) + 1)))
			for row in data.itertuples():
				store.append(list(self.to_string(x) for x in row))
			# Make the associated tree view
			view = Gtk.TreeView(store)
			columns = list(data.columns)
			columns.insert(0, data.index.name)
			for i, cname in enumerate(columns):
				renderer = Gtk.CellRendererText()
				column = Gtk.TreeViewColumn(cname, renderer, text=i)
				view.append_column(column)
			# Make a scrolled window containing the tree view
			scrolled = Gtk.ScrolledWindow()
			scrolled.set_hexpand(True)
			scrolled.set_vexpand(True)
			scrolled.add(view)
			self.add(scrolled)
			# Finalize window
			self.set_default_size(800, 600)

	def display(data, **kwargs):
		""" GTK-based display of a data frame.
		Args:
			data Data frame to display
			...	Forwarded keyword-arguments
		"""
		# Display given data
		gtk_run(lambda: _DataFrameDisplayWindow(data, **kwargs).show_all())
except Exception as err:
	def gtk_run(closure):
		""" Sink in case GTK cannot be used.
		Args:
			closure Ignored parameter
		"""
		tools.warning(f"GTK 3.0 is unavailable (in gtk_run): {err}")
	def display(data, **kwargs):
		""" Sink in case GTK cannot be used.
		Args:
			data Ignored
			...	Ignored keyword-arguments
		"""
		tools.warning(f"GTK 3.0 is unavailable (in display): {err}")

# ---------------------------------------------------------------------------- #
# Data frame columns selection helper

def select(data, *only_columns):
	""" "Intelligently" select columns from a data frame.
	Args:
		data Session or DataFrame to select
		...	Only columns to select, empty for all
	Returns:
		(Sub-)dataframe, by reference
	"""
	global Session
	# Unwrap data frame from session
	if isinstance(data, Session):
		data = data.data
	# Fast path
	if len(only_columns) == 0:
		return data
	# Intelligent selection
	columns = list()
	for only_column in only_columns:
		only_column = only_column.lower()
		for column in data.columns:
			if column not in columns and only_column in column.lower():
				columns.append(column)
	return data[columns]

# ---------------------------------------------------------------------------- #
# Training/evaluation data collection class

class Session:
	""" Training/evaluation data collection class.
	"""

	def __init__(self, path_results, location):
		""" Load the data from a training/evaluation result directory.
		Args:
			path_results Path-like to the result directory to load
		"""
		# Conversion to path
		if not isinstance(path_results, pathlib.Path):
			path_results = pathlib.Path(path_results)
		# Ensure directory exist
		if not path_results.exists():
			raise tools.UserException(f"Result directory {str(path_results)} cannot be accessed or does not exist")
		
		data_config = data_json = data_study = None
		# Load evaluation data
		#path_eval = path_results / "eval"
		path_eval = path_results / location
		try:
			data_eval = pandas.read_csv(path_eval, sep="\t", index_col=0)
			data_eval.index.name="Step number"
		except Exception as err:
			tools.warning(f"Result directory {str(path_results)}: unable to read evaluation data ({err})")
			data_eval = None
		# Merge data frames (if both are here)
		if data_study is not None and data_eval is not None:
			data = data_study.join(data_eval, how="outer")
		else:
			data = data_study or data_eval
		# Finalization
		self.name	 = path_results.name
		self.path	 = path_results
		self.config = data_config
		self.json	 = data_json
		self.data	 = data
		self.thresh = None

	def get(self, *only_columns):
		""" Get (some of) the data.
		Args:
			name Name of the data frame to consider
			...	Only columns to select, empty for all
		Returns:
			Selected data, by reference
		"""
		global select
		return select(self.data, *only_columns)

	def display(self, *only_columns, name=None):
		""" Just display (some of) the data.
		Args:
			name Name of the data frame to consider
			...	Only columns to select, empty for all
		Returns:
			self
		"""
		global display
		# Display the (selected sub)set
		display(self.get(*only_columns), title=(f"Session data{' (subset)' if len(only_columns) > 0 else ''} for {self.name!r}"))
		# Return self to enable chaining
		return self

	def compute_all(self):
		""" Carries all the automated computations.
		Returns:
			self
		"""
		# Carries all the computations
		for name, func in type(self).__dict__.items():
			if name == "compute_all":
				continue
			if name[:len("compute_")] == "compute_" and callable(func):
				func(self)
		# Return self to enable chaining
		return self

	def compute_epoch(self):
		""" Compute and append the epoch number, if not already done.
		Returns:
			self
		"""
		column_name = "Epoch number"
		# Check if already there
		if column_name in self.data.columns:
			return
		# Compute epoch number
		if self.json is None or "dataset" not in self.json:
			tools.warning("No valid JSON-formatted configuration, cannot compute the epoch number")
			return
		dataset_name	= self.json["dataset"]
		training_size = {"mnist": 60000, "fashionmnist": 60000, "cifar10": 50000, "cifar100": 50000}.get(dataset_name, None)
		if training_size is None:
			tools.warning(f"Unknown dataset {dataset_name!r}, cannot compute the epoch number")
			return
		self.data[column_name] = self.data["Training point count"] / training_size
		# Return self to enable chaining
		return self

	def compute_lr(self):
		""" Compute and append the learning rate, if not already done.
		Returns:
			self
		"""
		column_name = "Learning rate"
		# Check if already there
		if column_name in self.data.columns:
			return
		# Compute epoch number
		if self.json is None or "learning_rate" not in self.json:
			tools.warning("No valid JSON-formatted configuration, cannot compute the learning rate")
			return
		lr_schedule = self.json.get("learning_rate_schedule")
		if lr_schedule is None:
			lr = self.json["learning_rate"]
			lr_decay = self.json.get("learning_rate_decay", 0)
			lr_delta = self.json.get("learning_rate_decay_delta", 1)
			if lr_decay > 0:
				self.data[column_name] = lr / ((self.data.index // lr_delta * lr_delta) / lr_decay + 1)
			else:
				self.data[column_name] = lr
		else:
			tools.warning("Learning rate schedule not yet supported for schedule generation")
		# Return self to enable chaining
		return self

# ---------------------------------------------------------------------------- #
# Plot management class

class LinePlot:
	""" Line plot management class.
	"""
	# Plot style
	plt.style.use('fast')
	
	# Known line styles
	linestyles = ("-", "--", ":", "-.")

	# Known marker styles
	markstyles = ("+", "x", "1", "2", "3", "4")

	@classmethod
	def _get_line_style(self, ln):
		""" Get the line style and color for the given line number.
		Args:
			ln A non-negative integer representing the line number
		Returns:
			Associated line style, line color, marker style
		"""
		return self.linestyles[ln % len(self.linestyles)], f"C{ln}", self.markstyles[ln % len(self.markstyles)]

	def __init__(self, index=None):
		""" Title constructor.
		Args:
			index Column name to use as the index instead of the default
		"""
		# Make the subplots
		fig, ax = plt.subplots()
		# Store the non-finalized state
		self._fin = False # Not yet finalized
		self._fig = fig	 # Figure instance
		self._ax	= ax		# Original axis instance
		self._tax = None	# Twin axis instance
		self._axs = {}		# Map column names to axis (up to two)
		self._idx = index # Column name to use as index by default, None to use dataframe's index
		self._cnt = 0		 # Plot counter (to pick line style and color)

	def __del__(self):
		""" Close the figure on finalization.
		"""
		self.close()

	def _get_ax(self, name):
		""" Get the axis associated with the column selector, make it if possible.
		Args:
			name Column selector
		Returns:
			Associated axis
		"""
		# Return existing axis
		ax = self._axs.get(name, None)
		if ax is not None:
			return ax
		# Assert can make one more axis
		if len(self._axs) >= 2:
			raise RuntimeError("Line plot cannot have a 3rd y-axis")
		# Make one more axis
		if len(self._axs) == 0:
			ax = self._ax
		else:
			ax = self._ax.twinx()
			self._tax = ax
		self._axs[name] = ax
		# Return the axis
		return ax

	def include_simple(self, line, llbl, linewidth, mark=False, lalp=1., ccnt=None, post=None):
		""" Add the columns of the given data frame, can only be done before finalization.
		Args:
			line An iterable-once of (x, y) or (x, y, err)
			llbl Line label
			mark Whether to add point markers
			lalp Line alpha level
			ccnt Color and linestyle number to use
			post (Optional) post-operation lambda taking the axis object and returning None
		Returns:
			self
		"""
		# Assert not already finalized
		if self._fin:
			raise RuntimeError("Plot is already finalized and cannot include another line")
		# Prepare the data
		if isinstance(line, types.GeneratorType):
			line = tuple(line)
		if not isinstance(line, numpy.ndarray):
			line = numpy.array(line)
		if len(line.shape) != 2:
			raise RuntimeError(f"Unexpected line shape, expected (_, 2) or (_, 3), got {line.shape}")
		x = line[:,0]
		y = line[:,1]
		if line.shape[1] == 2:
			e = None
		elif line.shape[1] == 3:
			e = line[:,2]
		else:
			raise RuntimeError(f"Unexpected line shape, expected (_, 2) or (_, 3), got {line.shape}")
		# Recover the common, unique axis for the "simple" include
		axis = self._get_ax(None)
		# Pick a new line style and color
		linestyle, color, markstyle = self._get_line_style(self._cnt if ccnt is None else ccnt)
		# Plot the line and error line (if any)
		if e is not None:
			axis.fill_between(x, y - e, y + e, facecolor=color, alpha=0.2)
		axis.plot(x, y, marker=(markstyle if mark else None), label=llbl, linestyle=linestyle, color=color, alpha=lalp, linewidth=linewidth)
		# Optional post-operation on the axis
		if post is not None:
			post(axis)
		# Increase the counter only on success
		self._cnt += 1
		# Return self for chaining
		return self

	def include(self, data, *cols, errs=None, xticks= None, mark=None, linestyle = None, color = None, lalp=1., ccnt=None, xlogscale = False, logscale= False):
		""" Add the columns of the given data frame, can only be done before finalization.
		Args:
			data Session or dataframe holding the column(s) to add
			cols Column name(s) to include, mix selected columns together (same y-axis)
			errs Error suffix: for every selected column's real label, if a columns with 'real_label + errs' exists, it is used to display error bars
			mark Whether to add point markers
			lalp Line alpha level
			ccnt Color and linestyle number to use
		Returns:
			self
		"""
		#plt.style.use('ggplot')
		# Assert not already finalized
		if self._fin:
			raise RuntimeError("Plot is already finalized and cannot include another line")
		# Recover the dataframe if a session was given
		if isinstance(data, Session):
			data = data.data
		elif not isinstance(data, pandas.DataFrame):
			raise RuntimeError(f"Expected a Session or DataFrame for 'data', got a {tools.fullqual(type(data))!r}")
		# Get the x-axis values
		if xticks is not None:
			x = xticks
		elif self._idx is None:
			x = data.index.to_numpy()
		else:
			if self._idx not in data:
				raise RuntimeError(f"No column named {self._idx!r} to use as index in the given session/dataframe")
			x = data[self._idx].to_numpy()
		# Select semantic: empty list = select all
		if len(cols) == 0:
			cols = data.columns.to_list()
		# For every selection
		axis = None
		for col in cols:
			# Get associated data
			subd = select(data, col)
			# For every selected column
			for scol in subd:
				# Ignore index column
				if self._idx is not None and scol == self._idx:
					continue
				# Ignore error column
				if errs is not None and scol[:-len(errs)] in subd:
					continue
				# Get associated axis (if not done yet)
				if axis is None:
					axis = self._get_ax(col)
				# Pick a new line style and color
				#linestyle, color, markstyle = self._get_line_style(self._cnt if ccnt is None else ccnt)
				# Plot the data (line or error line)
				davg = subd[scol].to_numpy()
				errn = None if errs is None else (scol + errs)
				if errn is not None and errn in data:
					derr = data[errn].to_numpy()
					axis.fill_between(x, davg - derr, davg + derr, facecolor=color, alpha=0.2)
				axis.plot(x, davg, marker= mark, label=scol, linestyle=linestyle, color=color, alpha=lalp, linewidth=2.25)
				if logscale:
					axis.set_yscale('log')
				if xlogscale:
					axis.set_xscale('log')
				# Increase the counter only on success
				self._cnt += 1
			# Reset axis for next iteration
			axis = None
		# Return self for chaining
		return self

	def finalize(self, title, xlabel, ylabel, zlabel=None, xticks = None, tick_labels = None, xmin=None, xmax=None, ymin=None, ymax=None, zmin=None, zmax=None, legend=None):
		""" Finalize the plot, can be done only once and would prevent further inclusion.
		Args:
			title	Plot title
			xlabel Label for the x-axis
			ylabel Label for the y-axis
			zlabel Label for the twin y-axis, if any
			xmin	 Minimum for abscissa, if any
			xmax	 Maximum for abscissa, if any
			ymin	 Minimum for ordinate, if any
			ymax	 Maximum for ordinate, if any
			zmin	 Minimum for second ordinate, if any
			zmax	 Maximum for second ordinate, if any
			legend List of strings (one per 'include', in call order) to use as legend
		Returns:
			self
		"""
		# Fast path
		if self._fin:
			return self
		# Plot the legend
		def generator_sum(gen):
			res = None
			while True:
				try:
					val = next(gen)
					if res is None:
						res = val
					else:
						res += val
				except StopIteration:
					return res
		#(self._ax if self._tax is None else self._tax).legend(generator_sum(ax.get_legend_handles_labels()[0] for ax in self._axs.values()), generator_sum(ax.get_legend_handles_labels()[1] for ax in self._axs.values()) if legend is None else legend, loc="best", prop={'size': 13})
		# Style
		#plt.style.use('ggplot')
		# Legend
		if legend is not None:
			self._ax.legend(handles = legend, loc = 'upper left', bbox_to_anchor = (1, 1), fontsize = 10)
		# Plot the grid and labels
		self._ax.grid(which = 'major', axis = 'both')
		self._ax.set_xlabel(xlabel, fontsize=12)
		self._ax.set_ylabel(ylabel, fontsize=12)
		self._ax.set_title(title, fontsize=12)
		self._ax.tick_params(axis='y', which = 'both', labelsize=10)
		self._ax.tick_params(axis='x', which='major', labelsize=10)
		if zlabel is not None:
			if self._tax is None:
				tools.warning(f"No secondary y-axis found, but its label {zlabel!r} was provided")
			else:
				self._tax.set_ylabel(zlabel)
		elif self._tax is not None:
			tools.warning(f"No label provided for the secondary y-axis; using label {ylabel!r} from the primary")
			self._tax.set_ylabel(ylabel)
		if xticks is not None:
			self._ax.set_xticks(xticks)
		else:
			self._ax.set_xlim(left=xmin, right=xmax)
		if tick_labels is not None:
			self._ax.set_xticklabels(tick_labels)
		self._ax.set_ylim(bottom=ymin, top=ymax)
		if self._tax is not None:
			self._tax.set_ylim(bottom=zmin, top=zmax)
		# Mark finalized
		self._fin = True
		# Return self for chaining
		return self

	def display(self):
		""" Display the figure, which must have been finalized.
		Returns:
			self
		"""
		# Assert already finalized
		if not self._fin:
			raise RuntimeError("Cannot display a plot that has not been finalized yet")
		# Show the plot
		self._fig.show()
		# Return self for chaining
		return self

	def save(self, path, dpi=300, xsize=3, ysize=2):
		""" Save the figure, which must have been finalized.
		Args:
			path	Path of the file to write
			dpi	 Output image DPI (very good quality printing is usually 300 DPI)
			xsize Output image x-size (in cm)
			ysize Output image y-size (in cm)
		Returns:
			self
		"""
		# Assert already finalized
		if not self._fin:
			raise RuntimeError("Cannot display a plot that has not been finalized yet")
		# Save the figure
		self._fig.set_size_inches(xsize * 2.54, ysize * 2.54)
		self._fig.set_dpi(dpi)
		self._fig.savefig(path)
		# Return self for chaining
		return self

	def close(self):
		""" Explicitly "close" the associated figure (needed by pyplot), the instance cannot be used anymore after the call.
		"""
		if self._fig is not None: # The documentation of 'plt.close' does not explicitly specify that multiple calls are allowed on the same 'Figure'
			plt.close(self._fig)
			self._fig = None

class BoxPlot:
	""" Box/violin plot management class.
	"""

	def __init__(self, index=None):
		""" Title constructor.
		Args:
			index Column name to use as the index instead of the default
		"""
		# Make the subplots
		fig, ax = plt.subplots()
		# Store the non-finalized state
		self._fin	= False	# Not yet finalized
		self._fig	= fig		# Figure instance
		self._ax	 = ax		 # Original axis instance
		self._data = list() # Data: list of data array
		self._lbls = list() # Data: list of labels
		self._hls	= list() # List of horizontal line ordinates to plot

	def __del__(self):
		""" Close the figure on finalization.
		"""
		self.close()

	def include(self, data, label):
		""" Add the columns of the given data frame, can only be done before finalization.
		Args:
			data	Series or (numpy) array to add
			label Label for this data
		Returns:
			self
		"""
		# Assert not already finalized
		if self._fin:
			raise RuntimeError("Plot is already finalized and cannot include another line")
		# Recover the array if a series was given
		if isinstance(data, pandas.Series):
			data = data.to_numpy()
		elif not any(isinstance(data, dtype) for dtype in (numpy.ndarray, list, tuple)):
			raise RuntimeError(f"Expected a Series or an (numpy) array for 'data', got a {tools.fullqual(type(data))!r}")
		# Append the data
		self._data.append(data)
		self._lbls.append(label)
		# Return self for chaining
		return self

	def hline(self, y):
		""" Add an horizontal line to the plot.
		Args:
			y Ordinate of the horizontal line
		Returns:
			self
		"""
		# Push the ordinate
		self._hls.append(y)
		# Return self for chaining
		return self

	def finalize(self, title, ylabel, ymin=None, ymax=None, violin=False):
		""" Finalize the plot, can be done only once and would prevent further inclusion.
		Args:
			title	Plot title
			ylabel Label for the y-axis
			ymin	 Minimum for ordinate, if any
			ymax	 Maximum for ordinate, if any
			violin Whether to use violin plots instead of box plots
		Returns:
			self
		"""
		# Fast path
		if self._fin:
			return self
		# Plot the grid and labels
		self._ax.grid()
		self._ax.set_title(title)
		self._ax.set_ylabel(ylabel)
		self._ax.set_ylim(bottom=ymin, top=ymax)

		# Plot the data
		for i, y in enumerate(self._hls):
			self._ax.axhline(y, color=f"C{i}", linestyle="--")
		if violin:
			self._ax.violinplot(self._data, showmedians=True, showmeans=False, showextrema=False)
		else:
			self._ax.boxplot(self._data)
		plt.setp(self._ax, xticks=list(range(1, len(self._data) + 1)), xticklabels=self._lbls)

		# Mark finalized
		self._fin = True
		# Return self for chaining
		return self

	def display(self):
		""" Display the figure, which must have been finalized.
		Returns:
			self
		"""
		# Assert already finalized
		if not self._fin:
			raise RuntimeError("Cannot display a plot that has not been finalized yet")
		# Show the plot
		self._fig.show()
		# Return self for chaining
		return self

	def save(self, path, dpi=200, xsize=3, ysize=2):
		""" Save the figure, which must have been finalized.
		Args:
			path	Path of the file to write
			dpi	 Output image DPI (very good quality printing is usually 300 DPI)
			xsize Output image x-size (in cm)
			ysize Output image y-size (in cm)
		Returns:
			self
		"""
		# Assert already finalized
		if not self._fin:
			raise RuntimeError("Cannot display a plot that has not been finalized yet")
		# Save the figure
		self._fig.set_size_inches(xsize * 2.54, ysize * 2.54)
		self._fig.set_dpi(dpi)
		self._fig.savefig(path)
		# Return self for chaining
		return self

	def close(self):
		""" Explicitly "close" the associated figure (needed by pyplot), the instance cannot be used anymore after the call.
		"""
		if self._fig is not None: # The documentation of 'plt.close' does not explicitly specify that multiple calls are allowed on the same 'Figure'
			plt.close(self._fig)
			self._fig = None