IndianPlanetaryComputation/plotTest.py at master · meintte/IndianPlanetaryComputation · 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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
import numpy
import math
import matplotlib.animation as animation
import matplotlib.pyplot as plt
from matplotlib.font_manager import FontProperties

G = 6.67428e-11 # Newton's gravitational constant
AU = (149.6e6 * 1000) # Astronomical unit in metres
timestep = 24 * 3600 # One day in seconds (timestep for the simulation)

class planet(object):
    """
    Planet class.
    Contains information about the planet's
    current position, velocity and mass
    and other general information.
    """

    def __init__(self):
        self.px = 0.0
        self.py = 0.0
        self.vx = 0.0
        self.vy = 0.0
        self.mass = None
        self.color = None
        self.size = None
        self.name = None

    def compute_force(self, others):
        """
        Compute the total exerted force on the
        body at the current moment.
        """
        self.total_fx = self.total_fy = 0.0
        for other in others:
            # Compute the distance of the other body.
            sx, sy = self.px, self.py
            ox, oy = other.px, other.py
            dx = (ox-sx)
            dy = (oy-sy)
            d = numpy.sqrt(dx ** 2 + dy ** 2)

            # Compute the force of attraction
            f = G * self.mass * other.mass / (d ** 2)

            # Compute the direction of the force.
            theta = math.atan2(dy, dx)
            fx = math.cos(theta) * f
            fy = math.sin(theta) * f

            # Add to the total force exerted on the planet
            self.total_fx += fx
            self.total_fy += fy

    def update_position(self):
        """
        Update particle velocity and position based on the
        current exterted total force on the body.
        """
        self.vx += self.total_fx / self.mass * timestep
        self.vy += self.total_fy / self.mass * timestep
        self.px += self.vx * timestep
        self.py += self.vy * timestep

def animate(i, bodies, lines):
    """
    Animation function. Updates the
    plot on each interation.
    """
    for ind, body in enumerate(bodies):
        body.compute_force(numpy.delete(bodies, ind))
    for body in bodies:
        body.update_position()
    for i in range(len(bodies)):
        lines[i].set_data(bodies[i].px / AU, bodies[i].py / AU)
    return lines

def main():
    Sun = planet() # Instance of planet Sun
    Earth = planet() # Instance of planet Earth
    Venus = planet() # Instance of planet Venus

    Sun.mass = 1.98892 * 10 ** 30
    Sun.color = 'y'
    Sun.size = 50
    Sun.name = 'Sun'

    Earth.mass = 5.9742 * 10 ** 24
    Earth.px = -1 * AU
    Earth.vy = 29.783 * 1000
    Earth.color = 'b'
    Earth.size = 5
    Earth.name = 'Earth'

    Venus.mass = 4.8685 * 10 ** 24
    Venus.px = 0.723 * AU
    Venus.vy = -35.02 * 1000
    Venus.color = 'm'
    Venus.size = 5
    Venus.name = 'Venus'

    bodies = [Sun, Earth, Venus]
    lines = [None] * len(bodies)
    fig = plt.figure(figsize=(8,8))
    ax = plt.subplot()

    for i in range(len(bodies)):
        lines[i], = ax.plot(bodies[i].px / AU, bodies[i].py / AU,
        marker='o', color=bodies[i].color, ms=bodies[i].size,
        label=bodies[i].name)

    ani = animation.FuncAnimation(fig, animate, numpy.arange(1, 500),
        fargs=[bodies, lines], interval=20, blit=True, repeat=True)

    ax.set_xlabel('x [AU]')
    ax.set_ylabel('y [AU]')

    ax.set_xlim(-1.2, 1.2)
    ax.set_ylim(-1.2, 1.2)
    legend = ax.legend(loc=9, bbox_to_anchor=(0.5, 1.1), ncol=3)
    legend.legendHandles[0]._legmarker.set_markersize(6)

    circle = plt.Circle((0, 0), 0.8, fill=False)
    ax.add_artist(circle)

    plt.show()

if __name__ == "__main__":
    main()