Consistent colorscale for img_gen plots

[katiedagon]

It would be great to have the same colorbar/colorscale for the contour plots that are generated with img_gen=True. I think this would involve pre-setting a vmin and vmax to add to the plt.imshow() calls in the code below. The tricky bit would be determining those numbers which could vary depending on the run. It is nice to let the plotting code determine them for you, but then you lose consistency.

s2sml/applications/train.py

Lines 186 to 219 in f046727

	# save figs for later reference
	# selecting random sample in batch to visualize
	a_inp = img_noisy.cpu().detach().numpy()[:,nc-1,:,:]
	print(int(a_inp.shape[0]))
	samp_ = np.random.choice(np.arange(0,int(a_inp.shape[0]),1))
	# change plt logging level otherwise get a lot of debug output
	plt.set_loglevel(level='warning')
	# input
	im = plt.imshow(a_inp[samp_])
	plt.colorbar(im)
	plt.savefig(
	f"{save_loc}/trial{str(trial_num)}/{data_split}_inp_{str(epoch)}_{str(trial_num)}.png",
	bbox_inches='tight')
	plt.close()
	# output
	b_out = outputs.cpu().detach().numpy()[:,0,:,:][samp_]
	im = plt.imshow(b_out)
	plt.colorbar(im)
	plt.savefig(
	f"{save_loc}/trial{str(trial_num)}/{data_split}_out_{str(epoch)}_{str(trial_num)}.png",
	bbox_inches='tight')
	plt.close()
	# label
	c_lbl = img_label.cpu().detach().numpy()[:,0,:,:][samp_]
	im = plt.imshow(c_lbl)
	plt.colorbar(im)
	plt.savefig(
	f"{save_loc}/trial{str(trial_num)}/{data_split}_lbl_{str(epoch)}_{str(trial_num)}.png",
	bbox_inches='tight')
	plt.close()

I am happy to help work on this, adding an issue here to keep others in the loop.

[mariajmolina]

Thanks @katiedagon! Good idea! Based on discussion today, we can keep this as maybe lower priority to do item, to keep debugging / evaluation of output more modular, but I agree that it would be good to add an option for this as we are in the more final stages of training and saving models. Or maybe making this a hyper yaml file option that could be turned "on" or "off" could work to keep things more flexible when needed.

[katiedagon]

Sounds like a good plan @mariajmolina, thanks for clarifying that you already thought this through and in particular there are considerations that make setting a vmin and vmax tricky.