Extracted plots related stuff to own submodule

1 files changed, 38 insertions, 0 deletions

diff --git a/szilagyi/plots.py b/szilagyi/plots.py
new file mode 100644
index 0000000..e0727c8
--- /dev/null
+++ b/szilagyi/plots.py
@@ -0,0 +1,38 @@
+import matplotlib.pyplot as plot
+
+from .nomogram import *
+
+
+def grid(start, end, steps):
+	step = (end - start) / steps
+	i = start
+	while i <= end:
+		yield i
+		i += step
+
+
+if __name__ == "__main__":
+	swis = load("dataset")
+
+	def _swi(x, y):  # dt, depth
+		segments = find_boundary_curves(swis, x, y)
+		return calculate_swi(segments, x, y)
+
+	C = []
+	X = list(grid(0, 40, 1000))
+	Y = list(grid(0, 50000, 1000))
+	for y_scaled in range(0, 1000):
+		y = y_scaled * 50
+		row = []
+		for x_scaled in range(0, 1000):
+			x = x_scaled / 25
+			try:
+				swi = _swi(x, y)
+			except (IndexError, RuntimeError):
+				swi = -10
+			row.append(swi)
+		C.append(row)
+	plot.pcolormesh(X, Y, C, cmap='viridis', vmin=-10, vmax=10, rasterized=True)
+	for _, data in swis:
+		plot.plot([x[0] for x in data], [x[1] for x in data])
+	plot.show()