diff options
| -rw-r--r-- | szilagyi/plots.py | 37 |
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() |