View difference between Paste ID: gsnhZxB8 and 6dygw3vF
SHOW: | | - or go back to the newest paste.
1
#! /usr/bin/env python2
2
3
import numpy as np
4
from scipy.integrate import odeint
5
6
def Data_init(**kwargs):
7
  """
8
  if there is a kwarg['file'], then read from file, otherwise find
9
  kwargs['dimensions'] and kwargs['particles'] and randomize data.  If
10
  this fails, print an error.
11
12
  kwargs['time'] is a tuple of final time and number of steps.
13
  """
14
  if 'file' in kwargs:
15
    print "Reading the file"
16
  else:
17
    print "Randomizing the initial data"
18
    X = np.random.rand(kwargs['particles'],kwargs['dimensions']) * 2 - 1
19
    V = np.random.rand(kwargs['particles'],kwargs['dimensions']) * 2 - 1
20
    M = np.random.rand(kwargs['particles'])
21
22
  t_f,num = kwargs['time']
23
  t = np.linspace(0,t_f,num)
24
25
  return X,V,M,t
26
27
def Gravitational_force(X,M):
28
  # Expand the collections
29
  x1, x2 = X
30
  m1, m2 = M
31
32
  # Calculate a displacement vector
33
  r = x2-x1
34
  # Calculate it's magnitude
35
  r_mod = np.sqrt(r.dot(r))
36
  # Calculate the force itself
37
  f = m1*m2*r / np.power(r_mod,3)
38
39
  return f
40
41
def Net_force_1(X,M):
42-
  F = np.zeros(X.shape())
42+
  F = np.zeros(X.shape)
43-
  for i in xrange(X.shape()[0]):
43+
  for i in xrange(X.shape[0]):
44-
    for j in xrange(X.shape()[0]):
44+
    for j in xrange(X.shape[0]):
45
      if i == j:
46
        continue
47
      F += Gravitational_force(
48
          [ X[i,:], X[j,:] ],
49
          [ M[i], M[j] ]
50
          )
51
52
  return F
53
54
def ODE(W,t,M):
55
  """
56
  This is the function which is needed for odeint function
57
  """
58
  X, V = W
59
  F = Net_force_1(X,M)
60
  function = []
61-
  for i in xrange(X.shape()[0]):
61+
  for i in xrange(X.shape[0]):
62-
    for j in xrange(X.shape()[1]):
62+
    for j in xrange(X.shape[1]):
63
      function += [
64
          V[i,j],F[i,j]
65
          ]
66
  return function
67
68
# Initiate the data:
69
X0,V0,M,t = Data_init(particles=4,dimensions=2,time=[10,255])
70
71
# Solve the ODE system
72
solution = odeint(ODE,[X0,V0],t,args=(M,))