def aggregate_mass(self,too_close):
corrected_masses = AdaptingVectorQuantity()
corrected_positions = AdaptingVectorQuantity()
corrected_velocities = AdaptingVectorQuantity()
corrected_angular_momenta = AdaptingVectorQuantity()
for subset, m, pos, vel, Lin in zip(too_close, self.mass, self.position, self.velocity, self.angular_momentum):
if len(subset):
total_mass = subset.total_mass() + m
cmpos=(m*pos + subset.total_mass()*subset.center_of_mass())/total_mass
cmvel=(m*vel + subset.total_mass()*subset.center_of_mass_velocity())/total_mass
L=Lin+angular_momentum(m,pos-cmpos,vel-cmvel)+angular_momentum(subset.mass,subset.position-cmpos,subset.velocity-cmvel).sum(axis=0)
corrected_masses.append(total_mass)
corrected_positions.append(cmpos)
corrected_velocities.append(cmvel)
corrected_angular_momenta.append(L)
else:
corrected_masses.append(m)
corrected_positions.append(pos)
corrected_velocities.append(vel)
corrected_angular_momenta.append(Lin)
self.mass = corrected_masses
self.position = corrected_positions
self.velocity = corrected_velocities
self.angular_momentum = corrected_angular_momenta
def setup_point_positions(self):
self.x_positions = [
AdaptingVectorQuantity() for x in range(len(self.particles))
]
self.y_positions = [
AdaptingVectorQuantity() for x in range(len(self.particles))
]
def get_gravity_at_point(self,radius,x,y,z):
part=Particles(0)
for s in self.systems:
part.add_particles(s.particles)
ax=AdaptingVectorQuantity()
ay=AdaptingVectorQuantity()
az=AdaptingVectorQuantity()
for rr,xx,yy,zz in zip(radius,x,y,z):
dr2=((part.x-xx)**2+(part.y-yy)**2+(part.z-zz)**2+rr**2+part.radius**2)
ax.append( (self.G*part.mass*(part.x-xx)/dr2**1.5).sum() )
ay.append( (self.G*part.mass*(part.y-yy)/dr2**1.5).sum() )
az.append( (self.G*part.mass*(part.z-zz)/dr2**1.5).sum() )
return ax,ay,az
def create_N_vs_E(filename, N_list=None):
""" Iteratively runs run_hydrodynamics for several N in order to
create the data for the N vs. Energy error plot. Writes data to hdf5
file."""
if N_list:
N = N_list
else:
N = range(100, 1000, 50)
energy_errors = AdaptingVectorQuantity()
for nr_particles in N:
results = run_hydrodynamics(N=nr_particles, n_steps=50)
energy_begin = results.total_energy[0]
energy_end = results.total_energy[-1]
energy_error = (energy_begin - energy_end) / energy_end
energy_error = energy_error | units.no_unit
energy_errors.append(energy_error)
N_as_vq = N | units.no_unit
data = {'N': N_as_vq, 'energy_errors': energy_errors}
results = HydroResults(data)
write_to_hdf5('hydroresults/'+filename,\
results.__dict__)
pngfilename = options.N_vs_E.split('.')[0] + '.png'
plot_energy(results, filepath="plots/" + pngfilename)
def create_N_vs_t(filename, N_list=None):
""" Iteratively runs run_hydrodynamics for several N in order to
create the data for the N vs. wallclock-time plot. Writes data to
hdf5 file."""
if N_list:
N = N_list
else:
a = range(100, 1000, 100)
b = range(1000, 10000, 1000)
N = a + b
total_runtimes = AdaptingVectorQuantity()
for nr_particles in N:
start_time = time.time() | units.s
run_hydrodynamics(N=nr_particles, n_steps=1)
end_time = time.time() | units.s
total_runtime = end_time - start_time
total_runtimes.append(total_runtime)
N_as_vq = N | units.no_unit
data = {'N': N_as_vq, 'runtimes': total_runtimes}
results = HydroResults(data)
write_to_hdf5('hydroresults/'+filename,\
results.__dict__)
pngfilename = options.N_vs_t.split('.')[0] + '.png'
plot_steptimes(results, filepath="plots/" + pngfilename)
def get_potential_at_point(self,radius,x,y,z):
part=Particles(0)
for s in self.systems:
part.add_particles(s.particles)
phi=AdaptingVectorQuantity()
for rr,xx,yy,zz in zip(radius,x,y,z):
dr2=((part.x-xx)**2+(part.y-yy)**2+(part.z-zz)**2+rr**2+part.radius**2)
phi.append( (-self.G*part.mass/dr2**0.5).sum() )
return phi
def get_gravity_at_point(self, pos):
phi_0 = self.get_potential_at_point(pos)
grav = AdaptingVectorQuantity()
dpos = 0.001 * pos.length()
for ii in range(len(pos)):
ppos = 1.0 * pos
ppos[ii] += dpos
phi_1 = self.get_potential_at_point(ppos)
grav.append((phi_1 - phi_0) / dpos)
return grav
def run_pyth(interface, tend=100, dt=0.125, parameters=[]):
code = interface()
for name, value in parameters:
setattr(code.parameters, name, value)
code.particles.add_particles(new_particles())
code.commit_particles()
x = AdaptingVectorQuantity()
y = AdaptingVectorQuantity()
t = 0. | nbody_system.time
while (t < tend - dt / 2):
t = t + dt
code.evolve_model(t)
x.append(code.particles.x)
y.append(code.particles.y)
code.stop()
return x, y
def calculate_escape_time(index, number_of_systems):
numpy.random.seed()
# from amuse.community.smallN.muse_dynamics_mpi import SmallN
print "start of subprocess", index
x0 = AdaptingVectorQuantity()
y0 = AdaptingVectorQuantity()
x1 = AdaptingVectorQuantity()
y1 = AdaptingVectorQuantity()
x2 = AdaptingVectorQuantity()
y2 = AdaptingVectorQuantity()
m0 = AdaptingVectorQuantity()
m1 = AdaptingVectorQuantity()
m2 = AdaptingVectorQuantity()
tends = AdaptingVectorQuantity()
try:
for x in range(number_of_systems):
gravity = Hermite()
particles = new_initial_system()
code = SimulateTripleSystemUntilDecay(gravity, particles, index, x)
tend = code.evolve_model_until_escape()
print index, x, tend
code.stop()
x0.append(particles[0].x)
y0.append(particles[0].y)
x1.append(particles[1].x)
y1.append(particles[1].y)
x2.append(particles[2].x)
y2.append(particles[2].y)
m0.append(particles[0].mass)
m1.append(particles[1].mass)
m2.append(particles[2].mass)
tends.append(tend)
except KeyboardInterrupt:
pass
print "end of subprocess", index
return x0, y0, x1, y1, x2, y2, m0, m1, m2, tends
def vector(value=[], unit=None):
if unit is None:
if isinstance(value, core.unit):
return VectorQuantity([], unit=value)
elif isinstance(value, ScalarQuantity):
return value.as_vector_with_length(1)
else:
result = AdaptingVectorQuantity()
result.extend(value)
return result
else:
if isinstance(value, ScalarQuantity):
return value.as_vector_with_length(1)
else:
return VectorQuantity(value, unit)
def read_from_hdf5(filename):
""" Reads an hdf5 file and returns a HydroResult class. """
f = h5py.File(filename, 'r')
data = {}
for keyword in f.keys():
unitstring = f[keyword].attrs['unit']
evaluable_unit = parse_unitsstring(unitstring)
vq = f[keyword].value | eval(evaluable_unit)
avq = AdaptingVectorQuantity()
avq.extend(vq)
try:
mass_fractions = eval(f[keyword].attrs['mass_fractions'])
setattr(avq, 'mf', mass_fractions)
except KeyError, AttributeError:
pass
data.update({keyword: avq})
y2.append(particles[2].y)
m0.append(particles[0].mass)
m1.append(particles[1].mass)
m2.append(particles[2].mass)
tends.append(tend)
except KeyboardInterrupt:
pass
print "end of subprocess", index
return x0, y0, x1, y1, x2, y2, m0, m1, m2, tends
if __name__ == '__main__':
result = AdaptingVectorQuantity()
channel.MpiChannel.ensure_mpi_initialized()
rank = MPI.COMM_WORLD.Get_rank()
total_number_of_systems = int(sys.argv[1])
data = calculate_escape_time(rank, total_number_of_systems)
output = text.CsvFileText(filename="triple_data_{0:03}.csv".format(rank))
output.quantities = data
output.attribute_names = ("x0", "y0", "x1", "y1", "x2", "y2", "m0", "m1",
"m2", "tend")
output.store()
times = AdaptingVectorQuantity()
counts = AdaptingVectorQuantity()
def main():
dots = Dots()
line = Line()
f = h5py.File(args.filename, 'r')
fig = plt.figure()
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212)
for intr in f.values():
etas = []
final_smas_p0_vq = AdaptingVectorQuantity()
final_smas_p1_vq = AdaptingVectorQuantity()
for sim in intr.values():
etas.append(sim['eta'][0])
final_smas_p0_vq.append(sim['p0/sma'][-1]| retrieve_unit(sim['p0/sma']))
final_smas_p1_vq.append(sim['p1/sma'][-1]| retrieve_unit(sim['p1/sma']))
final_smas_p0 = final_smas_p0_vq.value_in(units.AU)
final_smas_p1 = final_smas_p1_vq.value_in(units.AU)
ax1.plot(etas, final_smas_p0, marker='o', label=intr.name, picker=5)
ax2.plot(etas, final_smas_p1, marker='o', label=intr.name, picker=5)
ax1.set_xlabel('eta')
ax2.set_xlabel('eta')
ax1.set_ylabel('final sma inner [AU]')
ax2.set_ylabel('final sma outer [AU]')
dset = f.values()[0].values()[0] #first dset available
time = quantify_dset(dset['time']).value_in(units.yr)
mass = quantify_dset(dset['mass']).value_in(units.MSun)
inner_sma0 = quantify_dset(dset['p0/sma']).value_in(units.AU)
outer_sma0 = quantify_dset(dset['p1/sma']).value_in(units.AU)
inner_sma_final_adiabatic_approx = sma_analytical(inner_sma0[0], 1.244e-5, time[-1], mass[0].sum() )
outer_sma_final_adiabatic_approx = sma_analytical(outer_sma0[0], 1.244e-5, time[-1], mass[0].sum() )
ax1.axhline(inner_sma_final_adiabatic_approx, xmin=0, xmax=1, c='m', label='adiabatic approx')
ax1.legend(loc='best')
ax2.axhline(outer_sma_final_adiabatic_approx, xmin=0, xmax=1, c='m', label='adiabatic approx')
ax2.legend(loc='best')
if args.logscale:
ax1.set_xscale('log')
ax2.set_xscale('log')
if args.ylim1:
ax1.set_ylim(args.ylim1)
if args.ylim2:
ax2.set_ylim(args.ylim2)
if args.xlim1:
ax1.set_xlim(args.xlim1)
if args.xlim2:
ax2.set_xlim(args.xlim2)
def onpick(event):
print("artist:{} ind:{}".format(event.artist, event.ind))
print("button: {}".format(event.mouseevent.button))
intr = f[event.artist.get_label()]
sim = intr.values()[event.ind]
time_vq = quantify_dset(sim['time'])
time = time_vq.value_in(units.yr)
position_vq = quantify_dset(sim['position'])
position = position_vq.value_in(units.AU)
CM_position_vq = quantify_dset(sim['CM_position'])
CM_position = CM_position_vq.value_in(units.AU)
x, y = 0, 1
central_x, central_y = position[:, 0, x] - CM_position[:,x], position[:, 0, y] - CM_position[:,y]
inner_x, inner_y = position[:, 1, x] - CM_position[:,x], position[:, 1, y] - CM_position[:,y]
outer_x, outer_y = position[:, 2, x] - CM_position[:,x], position[:, 2, y] - CM_position[:,y]
mass_vq = quantify_dset(sim['mass'])
mass = mass_vq[:,0].value_in(units.MSun)
if event.mouseevent.button == 1:
mu0 = quantify_dset(sim['mass'])[0].sum()
period_vq = quantify_dset(sim['p0/period'])
period = period_vq.value_in(units.yr)
true_anomaly = sim['p0/true_anomaly'].value
argument_of_periapsis = sim['p0/argument_of_periapsis'].value
sma_vq = quantify_dset(sim['p0/sma'])
sma = sma_vq.value_in(units.AU)
sma_an_vq = sma_analytical(sma_vq[0], 1.244e-5|(units.MSun/units.yr), time_vq, mu0)
sma_an = sma_an_vq.value_in(units.AU)
eccentricity = sim['p0/eccentricity'].value
newfig = plt.figure()
newax1 = newfig.add_subplot(511)
newax2 = newfig.add_subplot(512)
newax3 = newfig.add_subplot(513)
newax4 = newfig.add_subplot(514)
newax5 = newfig.add_subplot(515)
newax1.plot(time, sma, label='numerical')
newax1.plot(time, sma_an, label='analytical_adiabatic')
newax1.set_xlabel('time [yr]')
newax1.set_ylabel('sma [AU]')
newax1.legend(loc='best')
newax2.plot(time, eccentricity)
newax2.set_xlabel('time [yr]')
newax2.set_ylabel('eccentricity ')
newax3.plot(time, true_anomaly)
newax3.set_xlabel('time [yr]')
newax3.set_ylabel('true anomaly [degrees] ')
newax4.plot(time, argument_of_periapsis)
newax4.set_xlabel('time [yr]')
newax4.set_ylabel('argument of periapsis [degrees] ')
newax5.plot(time, period)
newax5.set_xlabel('time [yr]')
newax5.set_ylabel('period')
else:
newfig = plt.figure()
newax1 = newfig.add_subplot(321)
newax2 = newfig.add_subplot(322)
newax3 = newfig.add_subplot(323)
newax4 = newfig.add_subplot(324)
newax5 = newfig.add_subplot(325)
newax6 = newfig.add_subplot(326)
mass_vq = quantify_dset(sim['mass'])
mass = mass_vq.value_in(units.MSun)
kinetic_energy = sim['kinetic_energy'].value
potential_energy = sim['potential_energy'].value
total_energy = sim['total_energy'].value
CM_velocity_vq = quantify_dset(sim['CM_velocity'])
CM_velocity_mod = CM_velocity_vq.lengths().value_in(units.km/units.hour)
walltime = sim['walltime'].value - sim['walltime'][0]
newax1.plot(time, mass)
newax1.set_ylabel('central mass [MSun]')
newax1.set_xlabel('time [yr]')
newax2.plot(time, kinetic_energy, label='kinetic', **line.red)
newax2.plot(time, potential_energy,label='potential', **line.orange)
newax2.plot(time, total_energy, label='total', **line.white)
newax2.set_xlabel('time [yr]')
newax2.legend()
newax3.plot(central_x, central_y, **dots.white)
newax3.plot(inner_x, inner_y, **dots.red)
newax3.plot(outer_x, outer_y, **dots.yellow)
newax3.set_xlabel('x [AU]')
newax3.set_ylabel('y [AU]')
newax4.plot(CM_position[:, x], CM_position[:, y] )
newax4.set_xlim(-5, 5)
newax4.set_xlabel('CM position x [AU]')
newax4.set_ylabel('CM position y [AU]')
newax5.plot(time, CM_velocity_mod)
newax5.set_ylim(20, 30)
newax5.set_xlabel('time [yr]')
newax5.set_ylabel('CM velocity [km/hour]')
newax6.plot(time, walltime)
newax6.set_xlabel('time [yr]')
newax6.set_ylabel('walltime [s]')
newfig.show()
fig.canvas.mpl_connect('pick_event', onpick)
plt.show()
def test_run():
three = binary_with_planet(m1=0.6897 | units.MSun,
m2=0.20255 | units.MSun,
m_planet=0.333 | units.MJupiter,
r1=0.6489 | units.RSun,
r2=0.22623 | units.RSun,
r_planet=0.754 | units.RJupiter,
ecc_binary=0.15944,
P_binary=41.08 | units.day,
ecc_planet=0.00685,
a_planet=.7048 | units.AU,
pangle_planet=0.)
convert = nbody_system.nbody_to_si(1 | units.MSun, 1 | units.AU)
code = Huayno(convert)
code.parameters.inttype_parameter = code.inttypes.SHARED4
code.parameters.timestep_parameter = 0.1
# tend=100. | units.yr
tend = 100. | units.day
snapfreq = 1
dt = 10. | units.day
# dt=convert.to_si( 1. | nbody_system.time).in_(units.day)
code.particles.add_particles(three)
x = AdaptingVectorQuantity()
y = AdaptingVectorQuantity()
z = AdaptingVectorQuantity()
vx = AdaptingVectorQuantity()
vy = AdaptingVectorQuantity()
vz = AdaptingVectorQuantity()
x.append(code.particles.x)
y.append(code.particles.y)
z.append(code.particles.z)
vx.append(code.particles.vx)
vy.append(code.particles.vy)
vz.append(code.particles.vz)
ts = [0.]
E0 = code.potential_energy + code.kinetic_energy
dE = [1.e-14]
t = 0. | units.day
i = 0
while (t < tend - dt / 2):
i += 1
t = t + dt
if i % snapfreq == 0:
print t
ts.append(t.value_in(units.day))
code.evolve_model(t)
x.append(code.particles.x)
y.append(code.particles.y)
z.append(code.particles.z)
vx.append(code.particles.vx)
vy.append(code.particles.vy)
vz.append(code.particles.vz)
E = code.potential_energy + code.kinetic_energy
dE.append(abs(((E - E0) / E0)))
code.stop()
a, eps, pangle = elements(three.total_mass(), x[:, 2], y[:, 2], z[:, 2],
vx[:, 2], vy[:, 2], vz[:, 2])
x = x.value_in(units.AU)
y = y.value_in(units.AU)
a = a.value_in(units.AU)
eps = eps
print a[-1], eps[-1], pangle[-1]
f = pyplot.figure(figsize=(8, 8))
pyplot.plot(x[:, 0], y[:, 0], 'r.')
pyplot.plot(x[:, 1], y[:, 1], 'g.')
pyplot.plot(x[:, 2], y[:, 2], 'b.')
pyplot.xlim(-3, 3)
pyplot.ylim(-3, 3)
pyplot.xlabel('AU')
pyplot.savefig('three_16b.eps')
f = pyplot.figure(figsize=(8, 8))
pyplot.semilogy(ts, dE, 'g.')
pyplot.xlabel('time (day)')
pyplot.ylabel('dE/E0')
pyplot.savefig('three_16b_eerr.eps')
f = pyplot.figure(figsize=(8, 8))
pyplot.plot(ts, a, 'g.')
pyplot.xlabel('time (day)')
pyplot.ylabel('a (AU)')
pyplot.savefig('three_16b_a.eps')
f = pyplot.figure(figsize=(8, 8))
pyplot.plot(ts, eps, 'g.')
pyplot.xlabel('time (day)')
pyplot.ylabel('eccentricity')
pyplot.savefig('three_16b_ecc.eps')
f = pyplot.figure(figsize=(8, 8))
pyplot.plot(ts, pangle, 'g.')
pyplot.xlabel('time (day)')
pyplot.ylabel('long. of periapsis')
pyplot.savefig('three_16b_pangle.eps')
def run_hydrodynamics(N=100, Mtot=1|units.MSun, Rvir=1|units.RSun,
t_end=0.5|units.day, n_steps=10,\
vx = 0 |(units.RSun/units.day),\
vy = 0 |(units.RSun/units.day),\
vz = 0 |(units.RSun/units.day),\
plummer1=None, plummer2=None,\
bodyname = None):
""" Runs the hydrodynamics simulation and returns a HydroResults
instance.
FUNCTION WALKTHROUGH:
In the following explanation 'plummer1' and 'plummer2' are assumed
to be hdf5 files written by the function write_set_to_file().
Case 1:
If 'plummer1' and 'plummer2' are filenames of hdf5 files, then these
two plummer spheres will be smashed together.
Case 2:
If only plummer1 is supplied, then it will evolve plummer1 with
t_end timerange and n_steps steps.
Case 3:
If no plummers spheres are supplied, then it will generate a new
plummer sphere using the default/supplied initial conditions.
OUTPUT FILES:
If 'results_out' is specified, the HydroResult instance is written
to file in HDF5 format. This however does not use
write_set_to_file() which writes the entire Particles class and
its attributes to file at each dt, but uses write_to_hdf5()
from the 'support' module which is tailored to writing
HydroResults instances to file. This HDF5 contains all necessary
data to plot the required plots of the assignment.
In addition, the last snapshot of the Particles instance is written
to file using write_set_to_file(), the latter file is written to
the 'bodies' directory. Only the last snapshot is written to file
because the history of a Particle set is not of interest when
reloading them to smash plummer spheres. """
converter = nbody_system.nbody_to_si(Mtot, Rvir)
fi = Fi(converter)
if plummer1 and plummer2:
eta_smash = 0.3 | units.day
if plummer1 == plummer2:
bodies1 = read_set_from_file(plummer1, format='hdf5')
bodies2 = bodies1.copy()
bodies2.key += 1
else:
bodies1 = read_set_from_file(plummer1, format='hdf5')
bodies2 = read_set_from_file(plummer2, format='hdf5')
bodies1.move_to_center()
bodies2.move_to_center()
if vx.value_in(vx.unit) == 0 and vy.value_in(vy.unit) == 0 \
and vz.value_in(vz.unit) == 0:
bodies1.x += -1 | units.RSun
bodies2.x += 1 | units.RSun
else:
bodies1.x += (-1) * vx * eta_smash
bodies2.x += 1 * vx * eta_smash
bodies1.vx += vx
bodies2.vx += (-1) * vx
bodies1.vy += vy
bodies2.vy += (-1) * vy
bodies1.vz += vz
bodies2.vz += (-1) * vz
bodies1.add_particles(bodies2)
bodies = bodies1
elif plummer1 or plummer2:
if plummer1:
bodies = read_set_from_file(plummer1)
else:
bodies = read_set_from_file(plummer2)
bodies.move_to_center()
else:
bodies = new_plummer_gas_model(N, convert_nbody=converter)
fi.gas_particles.add_particles(bodies)
fi_to_framework = fi.gas_particles.new_channel_to(bodies)
fi.parameters.self_gravity_flag = True
data = {'lagrangianradii':AdaptingVectorQuantity(),\
'angular_momentum':AdaptingVectorQuantity(),\
'time':AdaptingVectorQuantity(),\
'positions':AdaptingVectorQuantity(),\
'kinetic_energy':AdaptingVectorQuantity(),\
'potential_energy':AdaptingVectorQuantity(),\
'total_energy':AdaptingVectorQuantity() }
mass_fractions = [0.10, 0.25, 0.50, 0.75]
setattr(data['lagrangianradii'], 'mf', mass_fractions)
data['radius_initial'], data['densities_initial'] = radial_density(\
fi.particles.x, fi.particles.mass, N=10, dim=1)
timerange = numpy.linspace(0, t_end.value_in(t_end.unit),\
n_steps) | t_end.unit
data['time'].extend(timerange)
fi.parameters.timestep = t_end / (n_steps + 1)
widget = drawwidget("Evolving")
pbar = pb.ProgressBar(widgets=widget, maxval=len(timerange)).start()
for i, t in enumerate(timerange):
fi.evolve_model(t)
data['kinetic_energy'].append(fi.kinetic_energy)
data['potential_energy'].append(fi.potential_energy)
data['total_energy'].append(fi.total_energy)
data['positions'].append(fi.particles.position)
data['angular_momentum'].append(fi.gas_particles.\
total_angular_momentum())
data['lagrangianradii'].append(fi.particles.LagrangianRadii(\
unit_converter=converter,\
mf=mass_fractions)[0])
if t == timerange[-1] and bodyname:
if os.path.dirname(bodyname) == '':
filename = "bodies/" + bodyname
fi_to_framework.copy()
write_set_to_file(bodies.savepoint(t), filename, "hdf5")
pbar.update(i)
pbar.finish()
data['radius_final'], data['densities_final'] = radial_density(\
fi.particles.x, fi.particles.mass, N=10, dim=1)
fi.stop()
results = HydroResults(data)
return results
