Cleaned up plots module

1 files changed, 25 insertions, 12 deletions

diff --git a/szilagyi/plots.py b/szilagyi/plots.py
index 9453919..a31f68a 100644
--- a/szilagyi/plots.py
+++ b/szilagyi/plots.py
@@ -1,32 +1,45 @@
 import os
+import sys
+import math
 
 import matplotlib.pyplot as plot
 
 from . import _dataset
-from .nomogram import *
+from . import calculate_swi
 
 
-def grid(start, end, steps):
+MAX_X = 40
+MAX_Y = 50000
+
+
+def _grid(start, end, steps):
 	step = (end - start) / steps
 	i = start
-	while i <= end:
+	while i < end or math.isclose(i, end):
 		yield i
 		i += step
 
 
-if __name__ == "__main__":
+def _scaled_range(steps, maximum):
+	step = maximum / steps
+	for i in range(steps):
+		yield i * step
+
+
+def colorful(width=1000, height=1000):
 	C = []
-	X = list(grid(0, 40, 1000))
-	Y = list(grid(0, 50000, 1000))
-	for y_scaled in range(0, 1000):
-		y = y_scaled * 50
+	X = list(_grid(0, MAX_X, width))
+	Y = list(_grid(0, MAX_Y, height))
+	for y in _scaled_range(height, MAX_Y):
 		row = []
-		for x_scaled in range(0, 1000):
-			x = x_scaled / 25
-			swi = calculate_swi(x, y)
-			row.append(swi)
+		for x in _scaled_range(width, MAX_X):
+			row.append(calculate_swi(x, y))
 		C.append(row)
 	plot.pcolormesh(X, Y, C, cmap='viridis', vmin=-10, vmax=10, rasterized=True)
 	for _, data in _dataset.INDICES:
 		plot.plot([x.x for x in data], [x.y for x in data])
 	plot.show()
+
+
+if __name__ == "__main__":
+	colorful()