nBodySystem/plot.py at master · ForrestHurley/nBodySystem · 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
import numpy as np
from matplotlib import pyplot as plt
import mpl_toolkits.mplot3d.axes3d as p3
from matplotlib import animation

class plotter(object):
    def __init__(self):
        self.x_range = (-2, 2)
        self.y_range = (-2, 2)
        self.z_range = (-2, 2)

        self.labels = ['X', 'Y', 'Z']
        self.title = "Orbits"

    def plot(self, data, show = True, set_axes = False): #data is bodies, times, position (x y z)
        self.data = data
        fig = plt.figure()
        ax = p3.Axes3D(fig)

        self.ln_plots = []

        for datum in self.data:
            self._draw_plot(ax, datum)

        if set_axes:
            ax.set_xlim3d(self.x_range)
            ax.set_ylim3d(self.y_range)
            ax.set_zlim3d(self.z_range)

        ax.set_xlabel(self.labels[0])
        ax.set_ylabel(self.labels[1])
        ax.set_zlabel(self.labels[2])

        ax.set_title(self.title)

        if show:
            plt.show()

        return fig

    def _draw_plot(self, ax, datum):
        self.ln_plots.append(ax.plot(datum[:,0],datum[:,1],datum[:,2]))

class anim_plotter(plotter):
    def __init__(self, rate = 1, history = 100, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.history = history
        self.rate = rate

    def update(self,n):
        start = max(n - self.history, 0)
        for body, datum in zip(self.pt_plots, self.data):
            body.set_data(datum[n][:2])
            body.set_3d_properties(datum[n][2])
        for body, datum in zip(self.ln_plots, self.data):
            body.set_data(datum[start:n,0],datum[start:n,1])
            body.set_3d_properties(datum[start:n,2])
        return self.ln_plots

    def plot(self, data, *args, **kwargs):
        kwargs['show'] = False

        self.pt_plots = []
        fig = super().plot(data, *args, **kwargs)

        count = self.data.shape[1]

        ani = animation.FuncAnimation(fig, self.update, count, interval = 50 / self.rate, blit = False)

        plt.show()
        return fig

    def _draw_plot(self, ax, datum):
        self.pt_plots.append(ax.plot([datum[0,0]],[datum[0,1]],[datum[0,2]],marker = 'o')[0])
        self.ln_plots.append(ax.plot([datum[0,0]],[datum[0,1]],[datum[0,2]])[0])