Draw ERT for all values and one value

draw_ert.py

+#encoding: utf-8
+from os import system, remove, path
+import matplotlib.pyplot as plt
+import numpy as np
+
+########################################################################
+# Interface
+########################################################################
+
+if __name__=="__main__":
+	import argparse
+
+	can = argparse.ArgumentParser()
+
+	can.add_argument("-N", "--nb-run", metavar="NB", default=100, type=int,
+			help="Number of runs")
+
+	can.add_argument("-v", "--for-value", metavar="NB", type=int,
+			help="Print ERT with probabilities to obtain this value")
+
+	arg = can.parse_known_args()
+	the = arg[0]
+
+	print(the)
+
+	# Minimum checks.
+	assert(0 < the.nb_run)
+
+	if path.isfile('run.csv'):
+		remove('run.csv')
+
+	for i in range(the.nb_run):
+		print("\rRun {}".format(i), end = '')
+		system("python3 snp.py -o run -no "+" ".join(arg[1]))
+	print("\r")
+
+	# Data extraction
+	data = []
+	attempt = {}
+	current_best = 0
+	temps_max = 0
+	value_max = 0
+	with open('run.csv', 'r') as fd:
+		for line in fd:
+			line = line.split(' ; ')
+			line[0] = int(line[0])
+			line[1] = float(line[1])
+
+			if line[0] > temps_max:
+				temps_max = line[0]
+
+			if line[1] > value_max:
+				value_max = line[1]
+
+			if line[0] == 1:
+				if len(attempt) > 0:
+					data.append(attempt)
+				current_best = float("inf")
+				attempt = {}
+			if line[1] < current_best:
+				attempt[line[0]] = line[1]
+				current_best = line[1]
+
+		data.append(attempt)
+
+	def probability(data,time,value):
+		num = 0
+		for attempt in data:
+			good = False
+			for t,v in attempt.items():
+				if t < time and v < value:
+					good = True
+					break
+			if good:
+				num += 1
+		return num/len(data)
+
+	probability = np.vectorize(probability, excluded=[0])
+
+	# Plot
+	if the.for_value is not None:
+		temps = np.linspace(0,temps_max,100)
+		P =  probability(data,temps,the.for_value)
+		plt.plot(temps,P)
+		plt.show()
+	else:
+		temps = np.linspace(0,temps_max,100)
+		values = np.linspace(0,value_max,100)
+		X, Y = np.meshgrid(temps, values)
+
+		Z = probability(data,X,Y)
+		fig, ax = plt.subplots()
+		cs = ax.contourf(X, Y, Z, 20)
+		cbar = fig.colorbar(cs)
+		plt.show()
+
+	remove('run.csv')