Extracted plots related stuff to own submodule

3 files changed, 39 insertions, 38 deletions

diff --git a/setup.py b/setup.py
index d488731..0a92236 100755
--- a/setup.py
+++ b/setup.py
@@ -5,4 +5,5 @@ from setuptools import setup
 setup(
 	name="szilagyi",
 	packages=["szilagyi"],
+	extra_requires={"plots": ["matplotlib"]},
 )
diff --git a/szilagyi/nomogram.py b/szilagyi/nomogram.py
index eb2665d..aed8dca 100644
--- a/szilagyi/nomogram.py
+++ b/szilagyi/nomogram.py
@@ -4,8 +4,6 @@ import os
 import re
 from collections import deque
 
-import matplotlib.pyplot as plot
-
 
 def load(directory):
 	def _read(iterable):
@@ -93,39 +91,3 @@ def calculate_swi(segments, x, y):
 	dist_to_low = min(dist(p[0], p[1], x, y) for p in (low[1][low[2]], low[1][low[2]]))
 	dist_to_high = min(dist(p[0], p[1], x, y) for p in (high[1][high[2]], high[1][high[2]]))
 	return dist_to_low / (dist_to_low + dist_to_high) * (high[0] - low[0]) + low[0]
-
-
-swis = load("dataset")
-
-
-def _swi(x, y):  # dt, depth
-	segments = find_boundary_curves(swis, x, y)
-	return calculate_swi(segments, x, y)
-
-
-def grid(start, end, steps):
-	step = (end - start) / steps
-	i = start
-	while i <= end:
-		yield i
-		i += step
-
-
-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()
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()