Penultimate Euler

"""
CMod Ex3: Program to run a 3D symplectic Euler time integration
of a particle in a gravitational potential. The program produces a plot
of the orbiting mass' trajectory and outputs, to a file, the
x and y positions of the mass at each time step.

Author: Taylor Baird
Version:26/10/16
"""
import sys
import matplotlib.pyplot as plt
from Particle3D import Particle3D
import numpy as np

G = 1.0                #Universal gravitational constant

#Check that an output file has been supplied
if len(sys.argv) != 2:
    print "Wrong number of arguments"
    print "Usage: " + sys.argv[0] + "<output file>"
    quit()
else:
    outfileName = sys.argv[1]

# open an output file for writing
outfile = open(outfileName, "w")

#Initialize particle with supplied mass and initial position and velocity.
mass = 0.0001
position = np.array([1,0,0])
velocity = np.array([0,1,0])
label = "orbiting mass"
p = Particle3D(mass, label, position, velocity)

# Initialize central particle. It is located at the origin, with a mass of 1.
#TODO Maybe remove this object since it is constant
mass = 1.0
label = "central_particle"
position = np.array([0, 0, 0])
velocity = np.array([0, 0, 0])
central_particle = Particle3D(mass, label, position, velocity)

#Set simulation parameters
numstep = 100
time = 0.0
timestep = 0.1

#Initialize x and y position lists and time list
tVals = [time]
xVals = [p.position[0]]
yVals = [p.position[1]]

# Calculation of the initial total energy of the particle as a reference
# so that we can calculate how the energy of the system fluctuates with successive timesteps.
potentialEnergy = -(G * central_particle.mass * p.mass) / np.linalg.norm(p.position)
initialTotalEnergy = potentialEnergy + p.getKineticEnergy()

outfile.write("{0:f}\t{1}\t{2}\n".format(time, p.position, initialTotalEnergy))

for i in range(numstep):
    #Use first order position update to update particle's position
    p.leapPos(timestep)
    #update force
    force = -((p.mass * central_particle.mass * G)/(np.linalg.norm(p.position))**3)*p.position
    #update particle velocity
    p.leapVel(timestep, force)

    #update time
    time += timestep
    #update time and position lists
    tVals.append(time)
    xVals.append(p.position[0])
    yVals.append(p.position[1])

    # Calculation of new total energy and fluctuation from initial total energy
    potentialEnergy = -(central_particle.mass * p.mass)/(np.linalg.norm(p.position))
    totalEnergy = potentialEnergy + p.getKineticEnergy()
    totalEnergyFluctuation = (totalEnergy - initialTotalEnergy)/initialTotalEnergy

    # TODO Format better, remove energy?
    # output current position and time to file
    outfile.write("{0:f}\t{1}\t{2}\t{3}\n".format(time, p.position, totalEnergy, totalEnergyFluctuation))

outfile.close()

plt.ylabel("y position")
plt.xlabel("x position")
plt.title("Plot of y position against x position for orbiting mass in\ngravitational potential using symplectic Euler method")
#Plot the y position of the particle against x position for successive time steps.
plt.plot(xVals, yVals)
#Make the scales of the axes equal
plt.gca().set_aspect('equal', adjustable='box')
#Save the obtained plot as a png image.
plt.savefig("symplecticEulerImg.png")
plt.show()