Commit c0de2e20 authored by Olaf's avatar Olaf

Save energies of all bodies

parent d489d565
......@@ -84,15 +84,22 @@ def calc(bodies, settings):
pass
else:
pass
positions, velocities, times, E_kin = dynamics(bodies, settings.dimensions, settings.timestep, settings.time, settings.barneshut, settings.theta, settings.approx, settings.timesave, settings.startsave)
positions, velocities, times, E_kin, E_pot = dynamics(bodies, settings.dimensions, settings.timestep,
settings.time, settings.barneshut, settings.theta,
settings.approx, settings.timesave, settings.startsave)
for i, body in enumerate(bodies):
if isinstance(body, Body):
body.add_trajectory((positions.transpose(1,0,2)[i]).tolist())
body.add_velocities((velocities[i]).tolist())
body.add_times(times)
body.add_kinetic_energy(E_kin.T[i])
body.add_potential_energy(E_pot.T[i])
elif isinstance(body, Asteroids):
body.add_trajectories((positions.transpose(1,0,2)[i:]).tolist())
body.add_velocities((velocities[i:]).tolist())
body.add_times(times)
body.add_kinetic_energy(E_kin.T[i:])
body.add_potential_energy(E_pot.T[i:])
# Fit 3D trajectory to 2D ellipse
# r, center, phi, axes = fit_ellipse(body.trajectory)
......
......@@ -328,5 +328,5 @@ def dynamics(bodies, D, h, t_max, barnes_hut, theta, approximation, j, startsave
all_pos_corr = np.transpose(np.transpose(all_pos, (1,0,2)) - np.transpose(all_pos, (1,0,2))[0], (1,0,2))
# return
return all_pos_corr, velo, times, E_kin
return all_pos_corr, velo, times, E_kin, E_pot
\ No newline at end of file
......@@ -155,6 +155,14 @@ class Body:
def add_ellfit_axes(self, axes):
"""Add axes of ellipsoid fitted on trajectory"""
self.ellfit_axes = axes
def add_kinetic_energy(self, Ekin):
"""Add kinetic energy at each timestep"""
self.E_kin = Ekin
def add_potential_energy(self, Epot):
"""Add potential energy at each timestep"""
self.E_pot = Epot
class Resonance_astroids:
"""Class with all information, settings and results of resonance astroid belt."""
......@@ -355,6 +363,18 @@ class Asteroids:
def add_velocities(self, velocities):
"""Add velocities of all asteroids"""
self.velocities = velocities
def add_kinetic_energy(self, Ekin):
"""Add kinetic energy at each timestep"""
self.E_kin = Ekin
def add_potential_energy(self, Epot):
"""Add potential energy at each timestep"""
self.E_pot = Epot
def add_times(self, times):
"""Add discrete time steps for which the trajectory is calculated"""
self.times = times
class Bodies:
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment