Add compute time function.

#matplotlib.rc("ytick", labelsize=14)

def compute_time(dataset, depths, widths):

	total = 0

	for d in depths:

		for w in widths:

			with open( os.path.join("../results", dataset, "depths1_{}_widths{}.json".format(d, w)), "r") as fp:

				results = json.load(fp)

			total += results["train"]["time"] + results["test"]["time"]

	return total

def heat_maps(dataset, depths, widths, save_csv=False):

	dir = os.path.join("figs", dataset, "hm")

	os.makedirs(dir, exist_ok=True)

	lc_depths = [1, 6, 12]

	lc_widths = [1, 32, 1024, 0]

	# generate figures

	for (dataset, dataset_name) in zip(datasets, dataset_names):

		if dataset == "boston":

			omit_iters = 0

		tex(dataset, dataset_name, lc_depths, lc_widths, lod_widths=hm_widths)

	# get total compute time of experiments

	total = 0

	for dataset in datasets:

		total += compute_time(dataset, hm_depths, hm_widths)

	print("Total compute time:", np.round(total/3600, 3), "hrs")

	print("Done!")

 No newline at end of file