Improve plotting script

.gitignore

@@ -3,3 +3,4 @@ legacy/
 __pycache__/
 *.pyc
 *.pyo
+*.png

plot_numberdealers.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 from datetime import timezone, datetime
-import parse_numberdealers
+from numberdealers import parse_numberdealers
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -10,8 +10,8 @@ def xy(messages):
 	x = []
 	y = []
 	for msg in messages:
-		x.append(datetime.fromtimestamp(msg.create_at / 1000, timezone.utc))
-		# x.append(msg.create_at / 1000)
+		# x.append(datetime.fromtimestamp(msg.create_at / 1000, timezone.utc))
+		x.append(msg.create_at / 1000)
 		y.append(msg.recognized_number)
 	return x, y
 
@@ -22,44 +22,71 @@ def find(xs, condition):
 			return i
 
 
-def abline(slope, intercept):
+def abline(slope, intercept, point1_x, point2_x):
 	"""Plot a line from slope and intercept"""
 	axes = plt.gca()
-	x_vals = np.array(axes.get_xlim())
+	x_vals = np.array([point1_x, point2_x])
 	y_vals = slope * x_vals + intercept
-	plt.plot(x_vals, y_vals, '--')
+	plt.plot(x_vals, y_vals, linewidth=1, color='grey')
 
 
 def main():
 	import sys
+	if len(sys.argv) <= 2:
+		print("Usage: plot_numberdealers.py <numberdealers.json> <numberdealers-ng.json>")
+		return 1
 	with open(sys.argv[1], encoding="utf-8") as fh:
 		numbers_og, _errors = parse_numberdealers.parse(fh)
 	with open(sys.argv[2], encoding="utf-8") as fh:
 		numbers_ng, _errors = parse_numberdealers.parse(fh)
 
-	# start_of_current_slope = find(numbers_og, lambda msg: msg.recognized_number <= 10464)
-	# slope_og = (
-		# numbers_og[-1].recognized_number - numbers_og[start_of_current_slope].recognized_number
-	# ) / (
-		# numbers_og[-1].create_at/1000 - numbers_og[start_of_current_slope].create_at/1000
-	# )
-	# # b = y - a x
-	# intercept_og = numbers_og[-1].recognized_number - slope_og * numbers_og[-1].create_at
+	start_of_current_slope = find(numbers_og, lambda msg: msg.create_at >= numbers_ng[0].create_at)
+	# a = (y2-y1) / (x2-x1)
+	slope_og = (
+		numbers_og[-1].recognized_number - numbers_og[start_of_current_slope].recognized_number
+	) / (
+		numbers_og[-1].create_at/1000 - numbers_og[start_of_current_slope].create_at/1000
+	)
+	# b = y - a x
+	intercept_og = numbers_og[-1].recognized_number - slope_og * numbers_og[-1].create_at/1000
 
-	# slope_ng = (
-		# numbers_ng[-1].recognized_number - numbers_ng[0].recognized_number
-	# ) / (
-		# numbers_ng[-1].create_at/1000 - numbers_ng[0].create_at/1000
-	# )
-	# intercept_ng = numbers_ng[-1].recognized_number - slope_ng * numbers_ng[-1].create_at
+	slope_ng = (
+		numbers_ng[-1].recognized_number - numbers_ng[0].recognized_number
+	) / (
+		numbers_ng[-1].create_at/1000 - numbers_ng[0].create_at/1000
+	)
+	intercept_ng = numbers_ng[-1].recognized_number - slope_ng * numbers_ng[-1].create_at/1000
 
-	fig, ax = plt.subplots()
+	# a1 x + b1 = a2 x + b2
+	# (a1 - a2) x = b2 - b1
+	# x = (b2 - b1) / (a1 - a2)
+	intercept_x = (intercept_ng - intercept_og) / (slope_og - slope_ng)
+	intercept_y = slope_ng * intercept_x + intercept_ng
+
+	fig, ax = plt.subplots(figsize=(15, 6))
 	ax.plot(*xy(numbers_og))
 	ax.plot(*xy(numbers_ng))
-	# abline(slope_og, -120000)
-	# abline(slope_ng, -500000)
-	plt.show()
+	abline(slope_og, intercept_og, numbers_og[-1].create_at/1000, intercept_x)
+	abline(slope_ng, intercept_ng, numbers_ng[-1].create_at/1000, intercept_x)
+
+	plt.margins(x=0)
+
+	locs = np.append(
+		np.linspace(numbers_og[0].create_at/1000, numbers_ng[0].create_at/1000, 6),
+		[numbers_ng[-1].create_at/1000, intercept_x]
+	)
+	# func = lambda x: x < intercept_x
+	# locs = np.append(locs[func(locs)][:-1], [intercept_x])
+	labels = ["{:%Y-%m-%d}".format(datetime.fromtimestamp(x, timezone.utc)) for x in locs]
+	plt.xticks(locs, labels)
+
+	plt.grid()
+
+	# plt.show()
+	plt.savefig("plot.png")
+	return 0
 
 
 if __name__ == "__main__":
-	main()
+	import sys
+	sys.exit(main())