FilesExpand file tree

 master

/

NBodySystem.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

100 lines (81 loc) · 3.37 KB

 master

/

NBodySystem.py

Copy path

Top

File metadata and controls

• Code
• Blame

100 lines (81 loc) · 3.37 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

from MySolarSystem import MySolarSystem

import numpy as np

import matplotlib.pyplot as plt

import seaborn as sns

class NBodySystem(MySolarSystem):

"""Docstring for NBodySystem. """

def __init__(self, seed):

"""TODO: to be defined1.

:seed: TODO

"""

MySolarSystem.__init__(self, seed)

self.data_filenames = 'data/{}NBody200yr{}'

self._seed = seed

def simulate_nbody(self, years = 8, planet_index = "", res = 5000):

self.G = 4*np.pi**2

n = res*years

numOfBodies = self.numberOfPlanets + 1

stop = years

dt = stop/float(n+1)

times = np.linspace(0,stop,n+1)

pos = np.zeros((n+1,2,numOfBodies))

vel = np.zeros_like(pos)

accel = np.zeros_like(pos)

self.star_velocity_init()

mass_array = np.append(self.mass,self.starMass)

x0 = np.append(self.x0, self.star_x0)

y0 = np.append(self.y0, self.star_y0)

vx0 = np.append(self.vx0, self.star_vx0)

vy0 = np.append(self.vy0, self.star_vy0)

pos[0] = np.array((x0, y0))

vel[0] = np.array((vx0, vy0))

accel[0] = self.accel_nbody(pos[0], mass_array)

print "Simulating %d years in %d timesteps"%(years, n)

for i in xrange(n):

if i%(n/100) == 0:

print 100*i/float(n)

pos[i+1] = pos[i] + vel[i]*dt + accel[i]*dt**2

accel[i+1] = self.accel_nbody(pos[i+1], mass_array)

vel[i+1] = vel[i] + 0.5*(accel[i]+accel[i+1])*dt

self.saveData(times, pos, vel, accel)

return pos, vel, accel, times

def star_velocity_init(self):

"""

Sets the velocity of the star so it matches the movement of the

planets (ensures a stationary center of mass for nbody-sim).

"""

starMass = self.starMass

masses = self.mass

self.star_vx0 = 0

self.star_vy0 = 0

self.star_x0 = 0

self.star_y0 = 0

for i in range(self.numberOfPlanets):

self.star_vx0 -= self.vx0[i]*masses[i]/starMass

self.star_vy0 -= self.vy0[i]*masses[i]/starMass

self.star_x0 -= self.x0[i]*masses[i]/starMass

self.star_y0 -= self.y0[i]*masses[i]/starMass

def loadData(self):

self.pos = np.load(self.data_filenames.format('pos','.npy'))

self.vel = np.load(self.data_filenames.format('vel','.npy'))

self.acc = np.load(self.data_filenames.format('acc','.npy'))

self.times = np.load(self.data_filenames.format('times','.npy'))

return self.pos, self.vel, self.acc, self.times

def saveData(self, times, pos, vel=None, acc=None):

np.save(self.data_filenames.format('pos',''), pos)

np.save(self.data_filenames.format('acc',''), acc)

np.save(self.data_filenames.format('vel',''), vel)

np.save(self.data_filenames.format('times',''), times)

def accel_nbody(self,pos,mass):

"""Used in simulate_nbody to find the acceleration of all the

bodies on eachother"""

N = len(mass)

f_array = np.zeros((N,2))

for i in range(N-1):

for j in range(i+1,N):

r_vec = pos[:,i] - pos[:,j]

r = np.linalg.norm(r_vec)

f = - self.G * mass[i]*mass[j] *r_vec/r**3

f_array[i] += f

f_array[j] -= f

return (f_array.T/mass)