GDPS/plotting.py at master · cbertinato/GDPS · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
import matplotlib.pyplot as plt
import matplotlib
from matplotlib.figure import Figure
from matplotlib.dates import DateFormatter, num2date
import matplotlib.patches as patches
import numpy as np

class Plots:
    def __init__(self):
        self.rects = []
        self.figure = None
        self.axes = []
        self.clicked = None

    def generate_subplots(self, x, *args):
        def _on_xlims_change(axes):
            # reset the x-axis format when the plot is resized
            axes.get_xaxis().set_major_formatter(DateFormatter('%H:%M:%S'))

        i = 0
        numplots = len(args)
        fig = plt.figure()

        self.cidclick = fig.canvas.mpl_connect('button_press_event', self.onclick)
        self.cidrelease = fig.canvas.mpl_connect('button_release_event', self.onrelease)
        self.cidmotion = fig.canvas.mpl_connect('motion_notify_event', self.onmotion)

        for arg in args:
            if i == 0:
                a = fig.add_subplot(numplots, 1, i+1)
            else:
                a = fig.add_subplot(numplots, 1, i+1, sharex=self.axes[0])

            a.plot(x.to_pydatetime(), arg)
            a.fmt_xdata = DateFormatter('%H:%M:%S')
            a.grid(True)
            a.callbacks.connect('xlim_changed', _on_xlims_change)
            self.axes.append(a)
            i += 1

        if not matplotlib.is_interactive():
            fig.show()

        self.figure = fig

    def onmotion(self, event):
        if self.clicked is None: return

        partners, index, x0, xclick, yclick = self.clicked

        # move rectangles
        dx = event.xdata - xclick
        for rect in partners:
            rect.set_x(x0 + dx)

        self.figure.canvas.draw()

    def onrelease(self, event):
        self.clicked = None
        self.figure.canvas.draw()

    def onclick(self, event):
        # TO DO: Don't place rectangle when zooming.
        # TO DO: Resize rectangles when plot extent changes.
        # TO DO: Make more efficient. Detect which plot is clicked?

        for partners in self.rects:
            index = 0
            for rect in partners:
                contains, attrd = rect.contains(event)
                if contains:
                    x0, _ = rect.xy
                    self.clicked = partners, index, x0, event.xdata, event.ydata
                    return
                index += 1

        partners = []
        for subplot in self.figure.axes:
            ylim = subplot.get_ylim()
            xlim = subplot.get_xlim()
            x_extent = (xlim[-1] - xlim[0]) * np.float64(0.1)

            # bottom left corner
            x0 = event.xdata - x_extent/2
            y0 = ylim[0]
            width = x_extent
            height = ylim[-1] - ylim[0]
            r = patches.Rectangle((x0, y0), width, height, alpha=0.1)
            # self.rects.append(r)
            partners.append(r)
            # rect = subplot.add_patch(r)
            subplot.add_patch(r)

        self.rects.append(partners)
        # self.rect.figure.canvas.draw()
        self.figure.canvas.draw()