import csv import math import operator import os import re from collections import deque import matplotlib.pyplot as plot def load(directory): def _read(iterable): for x, y in iterable: yield float(x), float(y) def _load(filename): with open(filename) as fd: reader = csv.reader(fd) return list(_read(reader)) def _files(directory): for file in os.listdir(directory): match = re.match(r"SWI_(-?\d+)\.csv", file) if match: yield int(match.group(1)), os.path.join(directory, file) return [(x, _load(y)) for x, y in sorted(_files(directory), key=lambda x: x[0])] def look_downwards(data, x, start): for i in range(start, 0, -1): if data[i - 1][0] < x: break else: raise IndexError return i - 1 def look_upwards(data, x, start): for i in range(start, len(data)): if data[i + 1][0] > x: break else: raise IndexError return i def find_segment(data, x): width = data[-1][0] - data[0][0] relative = x - data[0][0] candidate = math.floor(relative / width * len(data)) look = look_downwards if data[candidate][0] > x else look_upwards # May raise IndexError candidate = look(data, x, candidate) return candidate, candidate + 1 def find_boundary_curves(swis, x, y): segments = deque() for index, data in swis: i, j = find_segment(data, x) if data[i][1] > y and data[j][1] > y: segments.append((index, data, i, j)) break if data[i][1] < y and data[j][1] < y: if segments: segments.popleft() segments.append((index, data, i, j)) return segments swis = load("dataset") 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() fig, _ = plot.subplots() fig.canvas.mpl_connect('button_press_event', onclick) plot.show()