diff options
Diffstat (limited to 'plot.py')
| -rw-r--r-- | plot.py | 71 |
1 files changed, 58 insertions, 13 deletions
@@ -1,6 +1,5 @@ import csv import math -import operator import os import re from collections import deque @@ -65,22 +64,68 @@ def find_boundary_curves(swis, x, y): if segments: segments.popleft() segments.append((index, data, i, j)) + if len(segments) == 3: + middle = segments[1][1] + run = middle[j][0] - middle[i][0] + if run == 0: + raise RuntimeError # tidy up dataset + slope = (middle[j][1] - middle[i][1]) / run + intercept = middle[j][1] - slope * middle[j][0] + value = slope * x + intercept + if value == y: + raise RuntimeError # Exactly on point; SWI == index + if value < y: + segments.popleft() + else: + segments.pop() + if len(segments) == 1: + raise RuntimeError # SWI == -10 return segments -swis = load("dataset") +def dist(x1, y1, x2, y2): + return math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) + + +def calculate_swi(segments, x, y): + low = segments[0] + high = segments[1] + 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] -def onclick(event): - if event.button != 1: - return - plot.clf() - segments = find_boundary_curves(swis, event.xdata, event.ydata) - plot.plot([event.xdata], [event.ydata], "rx") - for index, data, i, j in segments: - plot.plot([x[0] for x in data], [x[1] for x in data], ".", label=index) - plot.show() +swis = load("dataset") + -fig, _ = plot.subplots() -fig.canvas.mpl_connect('button_press_event', onclick) +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() |