def test19(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test19"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
stars = Particles(2)
stars[0].x = 1.0 | units.km
stars[1].x = 2.0 | units.km
binaries = Particles(1)
binaries[0].y = 1.0 | units.km
binaries[0].child1 = stars[0]
binaries[0].child2 = stars[1]
stars[0].parent = binaries[0]
stars[1].parent = binaries[0]
self.assertEqual(binaries[0].child1,stars[0])
self.assertEqual(binaries[0].child2,stars[1])
self.assertEqual(binaries[0].child1.parent,binaries[0])
self.assertEqual(binaries[0].child2.parent,binaries[0])
io.write_set_to_file([binaries,stars], output_file, "hdf5", names = ['binaries', 'children'], version = self.store_version())
loader_binaries, loaded_stars = io.read_set_from_file(output_file, "hdf5", names = ['binaries', 'children'], version = self.store_version())
self.assertEqual(loader_binaries[0].child1.key,stars[0].key)
self.assertEqual(loader_binaries[0].child2.key,stars[1].key)
self.assertEqual(loader_binaries[0].child1,loaded_stars[0])
self.assertEqual(loader_binaries[0].child2,loaded_stars[1])
self.assertEqual(loader_binaries[0].child1.parent,loader_binaries[0])
self.assertEqual(loader_binaries[0].child2.parent,loader_binaries[0])
def test9(self):
p = datamodel.Particles(5)
p.a = [1.0, 2.0, 3.0, 4.0, 5.0]
p.b = [10, 11, 12, 13, 14] | units.m
p.c = [20, 21, 22, 23, 24] | units.m
io.write_set_to_file(
p,
"test.csv",
"txt",
attribute_names = ('a', 'b', 'c'),
attribute_types = (None, units.m, units.m),
maximum_number_of_lines_buffered = 1,
)
with open("test.csv", "r") as f:
contents = f.read()
expected_contents = '#a b c\n#- m m\n1.0 10.0 20.0\n2.0 11.0 21.0\n3.0 12.0 22.0\n4.0 13.0 23.0\n5.0 14.0 24.0\n'
self.assertEquals(expected_contents, contents)
p2 = io.read_set_from_file(
"test.csv",
"txt",
attribute_names = ('a', 'b', 'c'),
attribute_types = (None, units.m, units.m),
maximum_number_of_lines_buffered = 1,
)
self.assertAlmostRelativeEquals(p2.a, p.a)
self.assertAlmostRelativeEquals(p2.b, p.b)
self.assertAlmostRelativeEquals(p2.c, p.c)
def test11(self):
directory_name = os.path.dirname(__file__)
filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
gas, halo, disk, bulge, stars, bndry = io.read_set_from_file(filename, format='gadget')
self.assertEquals(len(gas), 1472)
self.assertEquals(len(halo), 0)
self.assertEquals(gas[0].key,1)
self.assertEquals(gas[1].key,2)
self.assertEquals(gas[2].key,3)
self.assertEquals(gas[1471].key,1472)
self.assertAlmostRelativeEquals(gas[0:5].x,[-0.0713372901082, 0.0713372901082, -0.21178227663, -0.0698266476393, 0.0698266476393] | nbody_system.length, 7)
self.assertAlmostRelativeEquals(gas[0:5].u, [0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451] | (nbody_system.length / nbody_system.time)**2, 7 )
outputfilename = 'gadgettest.output'
try:
io.write_set_to_file((gas, halo, disk, bulge, stars, bndry), outputfilename, format='gadget', ids_are_long = False)
gas, halo, disk, bulge, stars, bndry = io.read_set_from_file(outputfilename, format='gadget')
self.assertEquals(len(gas), 1472)
self.assertEquals(len(halo), 0)
self.assertEquals(gas[0].key,1)
self.assertEquals(gas[1].key,2)
self.assertEquals(gas[2].key,3)
self.assertEquals(gas[1471].key,1472)
finally:
if os.path.exists(outputfilename):
os.remove(outputfilename)
def test9(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test9"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
number_of_particles = 10
p = Particles(number_of_particles)
p.mass = [x * 2.0 for x in range(number_of_particles)] | units.kg
p.model_time = 2.0 | units.s
io.write_set_to_file(
p,
output_file,
"hdf5",
timestamp = 2 | units.Myr,
scale = 1 | units.kg,
version = self.store_version(),
close_file = True
)
loaded_particles = io.read_set_from_file(
output_file,
"hdf5",
version = self.store_version()
)
loaded_particles = self.get_version_in_store(loaded_particles)
a = loaded_particles.collection_attributes
self.assertAlmostRelativeEquals(a.timestamp, 2 | units.Myr)
self.assertAlmostRelativeEquals(a.scale, 1 | units.kg)
def test10(self):
p = datamodel.Particles(keys=[30,31,32,33,34])
p.a = [1.0, 2.0, 3.0, 4.0, 5.0]
p.b = [10, 11, 12, 13, 14] | units.m
p.c = [20, 21, 22, 23, 24] | units.m
print p.key
io.write_set_to_file(
p,
"test.csv",
"txt",
attribute_names = ('a', 'b', 'c'),
attribute_types = (None, units.m, units.m),
maximum_number_of_lines_buffered = 1,
key_in_column = 0
)
with open("test.csv", "r") as f:
contents = f.read()
print repr(contents)
expected_contents = '#a b c\n#- m m\n30 1.0 10.0 20.0\n31 2.0 11.0 21.0\n32 3.0 12.0 22.0\n33 4.0 13.0 23.0\n34 5.0 14.0 24.0\n'
self.assertEquals(expected_contents, contents)
p2 = io.read_set_from_file(
"test.csv",
"txt",
attribute_names = ('a', 'b', 'c'),
attribute_types = (None, units.m, units.m),
maximum_number_of_lines_buffered = 1,
key_in_column = 0
)
self.assertEquals(p2.key, p.key)
self.assertAlmostRelativeEquals(p2.a, p.a)
self.assertAlmostRelativeEquals(p2.b, p.b)
self.assertAlmostRelativeEquals(p2.c, p.c)
def hybrid_self_gravitating_cluster(Ncl, Mstar, Rcl, t_end, dt,
Mcut, filename):
""" Create and evolve hybrid cluster; store data to file """
stars_superset, stars_below_cut, stars_above_cut, nbody_converter =\
create_cluster(Ncl, Mstar, Rcl, Mcut)
bridge, channels = setup_codes(stars_below_cut, stars_above_cut,
nbody_converter)
time = 0. | units.Myr
write_set_to_file(stars_superset.savepoint(time), filename, 'amuse',
append_to_file=False)
print "Saving data to", filename
while time < (t_end - dt/2.):
time += dt
print "Evolving to", time
# NB this is the solver, not the bridge when only one solver is used!
# evolve_model, however, can be called both on bridge and solver (:
bridge.evolve_model(time)
channels.copy()
write_set_to_file(stars_superset.savepoint(time),
filename, 'amuse')
def main(
star_mass = 1|units.MSun,
star_radius = 1|units.RSun,
disk_minumum_radius = 0.05 | units.AU,
disk_maximum_radius = 10 | units.AU,
disk_mass = 200 | MEarth,
accurancy = 0.0001,
planet_density = 3 | units.g/units.cm**3,
output_filename = 'star.h5',
seed = -1):
if seed < 0:
rng = random.RandomState()
else:
rng = random.RandomState(seed)
output = new_system(
star_mass,
star_radius,
disk_minumum_radius,
disk_maximum_radius,
disk_mass,
accurancy,
planet_density,
rng = rng)
star = output[0]
print "Number of planets generated:", len(star.planets)
print "Total mass:", star.planets.mass.sum().as_quantity_in(MEarth)
for i, planet in enumerate(star.planets):
print "Planet: {0: 3d} , mass: {1: 8.3f} MEarth, a: {2: 8.2f} AU".format(i, planet.mass.value_in(MEarth), planet.semimajor_axis.value_in(units.AU))
write_set_to_file(output, output_filename, 'hdf5', version="2.0",append_to_file = False)
def test11(self):
print "Testing saving/loading timestamp in NEMO"
x = datamodel.Particles(2)
convert = nbody_system.nbody_to_si(1 | units.kg, 2 | units.m)
x.mass = [1.0, 2.0] | units.kg
x.position = [[1,2,3], [3,5,6]] | units.m
x.velocity = [[1,2,3], [3,5,6]] | units.m / units.s
current_time = 1.0 | units.Myr
io.write_set_to_file(x.savepoint(current_time), "time_test_unit.tsf","tsf", nbody_to_si_converter = convert)
y = io.read_set_from_file("time_test_unit.tsf","tsf", nbody_to_si_converter = convert)
self.assertAlmostEquals(current_time, y.previous_state().get_timestamp(), 8, in_units=units.Myr)
self.assertAlmostEquals(x.mass, y.mass, 8)
self.assertAlmostEquals(x.position, y.position,8)
self.assertAlmostEquals(x.velocity, y.velocity,8)
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | nbody_system.mass
x.position = [[1,2,3], [3,5,6]] | nbody_system.length
x.velocity = [[1,2,3], [3,5,6]] | nbody_system.speed
current_time = 1.0 | nbody_system.time
io.write_set_to_file(x.savepoint(current_time), "time_test_unit.tsf","tsf")
y = io.read_set_from_file("time_test_unit.tsf","tsf")
self.assertAlmostEquals(current_time, y.previous_state().get_timestamp(), 8, in_units=nbody_system.time)
self.assertAlmostEquals(x.mass, y.mass, 8)
self.assertAlmostEquals(x.position, y.position,8)
self.assertAlmostEquals(x.velocity, y.velocity,8)
os.remove("time_test_unit.tsf")
def make_galaxies():
if os.path.exists('disk_galactICs.amuse'):
galaxy1 = read_set_from_file('disk_galactICs.amuse', 'amuse')
else:
halo_number_of_particles = NHALO
converter = nbody_system.nbody_to_si(
1.0e9 | units.MSun, 1. | units.kpc)
galaxy1 = new_galactics_model(
halo_number_of_particles,
disk_number_of_particles=NDISK,
generate_bulge_flag=False,
unit_system_converter=converter,
disk_random_seed=12345)
write_set_to_file(galaxy1, 'disk_galactICs.amuse', 'amuse')
galaxy2 = Particles(len(galaxy1))
galaxy2.mass = galaxy1.mass
galaxy2.position = galaxy1.position
galaxy2.velocity = galaxy1.velocity
galaxy1.rotate(0.0, numpy.pi/4, 0.0)
galaxy2.rotate(numpy.pi/6, 0.0, 0.0)
galaxy1.position += [100.0, 0, 0] | units.kpc
galaxy2.position -= [100.0, 0, 0] | units.kpc
galaxy1.velocity += [0.0, 50.0, 0] | units.km/units.s
galaxy2.velocity -= [0.0, 50.0, 0] | units.km/units.s
return galaxy1, galaxy2
def test5(self):
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | units.MSun
x.radius = [3.0, 4.0] | units.RSun
expected = [
"#mass radius\n#MSun RSun\n{0} 1.0 3.0\n{1} 2.0 4.0\n".format(x[0].key, x[1].key),
"#mass radius\n#MSun RSun\n1.0 {0} 3.0\n2.0 {1} 4.0\n".format(x[0].key, x[1].key),
"#mass radius\n#MSun RSun\n1.0 3.0 {0}\n2.0 4.0 {1}\n".format(x[0].key, x[1].key),
]
for column_index, expected_content in enumerate(expected):
io.write_set_to_file(
x,
"test.csv",
"txt",
key_in_column = column_index,
attribute_types = (units.MSun, units.RSun)
)
with open("test.csv", "r") as f:
contents = f.read()
self.assertEquals(expected_content, contents)
y = io.read_set_from_file(
"test.csv",
"txt",
key_in_column = column_index,
attribute_types = (units.MSun, units.RSun),
attribute_names = ('mass', 'radius')
)
self.assertEquals(y[0], x[0])
self.assertEquals(y[1], x[1])
os.remove("test.csv")
def test58(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test58"+self.store_version()+".h5")
if os.path.exists(output_file):
os.remove(output_file)
x = Particles(keys = (1,2))
x.mass = [1,2] | units.kg
y = Particles(keys = (11,12,13,14))
y.mass = [40,50,60,70] | units.kg
x.sub = None
x[0].sub = y[0:2]
y[0].sub = y[2:]
io.write_set_to_file(x, output_file,"amuse", version=self.store_version())
z = io.read_set_from_file(output_file,"amuse")
self.assertEqual(len(x), len(z))
self.assertEqual(len(x[0].sub), len(z[0].sub))
self.assertEqual(x[0].sub[0].sub.key, (13,14))
self.assertEqual(z[0].sub[0].sub.key, (13,14))
self.assertEqual(id(x[0].sub._original_set()), id(x[0].sub[0].sub._original_set()))
# no more subsets, could this be fixed, would be very nice to do so
# self.assertEqual(id(z[0].sub._original_set()), id(z[0].sub[0].sub._original_set()))
self.assertTrue(z[1].sub==x[1].sub==None)
os.remove(output_file)
def test16(self):
import h5py
print h5py.version.version
print h5py
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test16"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
stars = Particles(2)
stars[0].x = 1.0 | units.km
stars[1].x = 2.0 | units.km
stars[0].nn = stars[1]
stars[1].nn = stars[0]
self.assertEqual(stars[0].nn,stars[1])
self.assertEqual(stars[1].nn,stars[0])
io.write_set_to_file(stars, output_file, "hdf5", version = self.store_version())
processor = self.store_factory()(output_file, True, open_for_writing = True)
loaded = processor.load()
self.assertEqual(loaded[0].nn,loaded[1])
self.assertEqual(loaded[1].nn,loaded[0])
self.assertEqual(loaded[0].nn.key,stars[1].key)
self.assertEqual(loaded[0].nn.x,stars[1].x)
def test17(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test17"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
stars = Particles(4)
stars[0].x = 1.0 | units.km
stars[1].x = 2.0 | units.km
stars[2].x = 3.0 | units.km
stars[3].x = 4.0 | units.km
binaries = Particles(2)
binaries[0].y = 1.0 | units.km
binaries[1].y = 2.0 | units.km
binaries[0].child1 = stars[0]
binaries[0].child2 = stars[1]
binaries[1].child1 = stars[2]
binaries[1].child2 = stars[3]
self.assertEqual(binaries[0].child1,stars[0])
self.assertEqual(binaries[1].child1,stars[2])
io.write_set_to_file(binaries, output_file, "hdf5", version = self.store_version())
loaded = io.read_set_from_file(output_file, "hdf5", version = self.store_version())
self.assertEqual(loaded[0].child1.key,stars[0].key)
self.assertEqual(loaded[1].child1.key,stars[2].key)
def write_output(filename, parts, conv):
particles_nbody = ParticlesWithUnitsConverted(parts, conv.as_converter_from_nbody_to_si())
#print particles_nbody
write_set_to_file(particles_nbody, filename, "txt", attribute_names= ('rho', 'mass', 'x', 'y', 'z','vx', 'vy', 'vz'))
return 0
def save_particles_to_file(bodies, bodies_to_save, bodies_filename,time,converter):
#Add attributes that I'm interested in to the bodies_to_save
bodies_to_save.position = converter.to_si(bodies.position).as_quantity_in(units.AU)
bodies_to_save.velocity = converter.to_si(bodies.velocity).as_quantity_in(units.kms)
bodies_to_save.semimajoraxis = converter.to_si(bodies.semimajoraxis).as_quantity_in(units.AU)
bodies_to_save.eccentricity = bodies.eccentricity
bodies_to_save.time = converter.to_si(time).as_quantity_in(units.yr)
write_set_to_file(bodies_to_save, bodies_filename, "hdf5")
def test3(self):
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | units.MSun
x.radius = [3.0, 4.0] | units.RSun
io.write_set_to_file(x, "test.csv","txt", attribute_types = (units.MSun, units.RSun))
with open("test.csv", "r") as f:
contents = f.read()
self.assertEquals("#mass radius\n#MSun RSun\n1.0 3.0\n2.0 4.0\n", contents)
os.remove("test.csv")
def evolve_disk_flyby(bodies, gravity,
t_end, n_steps, converter, snap_dir, file_out):
bodies = ParticlesSuperset([stars, planetesimals])
channel_from_gr_to_framework = gravity.particles.new_channel_to(bodies)
rm_file(file_out)
Etot_init = gravity.kinetic_energy + gravity.potential_energy
Etot = Etot_init
dt = t_end / float(n_steps)
time = 0.0 | units.yr
stdout = (file_out.split('.'))[0]
stdout += '.txt'
f = open(snap_dir + "/" + stdout, 'a')
print " ** evolving: t_end = ", t_end, ", dt = ", dt.in_(units.yr)
print " \t\t", "time", "\t\t", "dE"
while time<=t_end:
gravity.evolve_model(time)
channel_from_gr_to_framework.copy()
bodies.collection_attributes.timestamp = time
Ekin = gravity.kinetic_energy
Epot = gravity.potential_energy
Etot = Ekin + Epot
dE = Etot_init-Etot
nb_E = converter.to_nbody(Etot)
#nb_J = converter.nbody_length ** 2 * units.nbody_mass * units.nbody_time ** -2 # not supposed to work
# A formatted string would work better than tabs. (Tabs never work)
line = " \t" + str(time.value_in(units.yr)) + "\t" + str(nb_E) + "\t" + str(dE/Etot_init)
print line
f.write(line + "\n")
# Write coordinates in Center of Mass frame
# Move Stars to CoM (initially in CoM)
#### The stars are already in CoM coordinates ####
# Move planetesimals to CoM (initially in CoM)
#### The stars are already in CoM coordinates #### (Does this mean the other method is wrong?)
write_set_to_file(bodies, snap_dir + "/" + file_out, "hdf5")
time += dt
gravity.stop()
return
def store_grids(self, grids, step):
if __name__ == "__plot__":
return
grids_in_memory = [x.copy() for x in grids]
io.write_set_to_file(
grids_in_memory,
"kelvin_helmholtz_{2}_{0}_{1}.vtu".format(self.number_of_grid_points, step, self.name_of_the_code),
"vtu",
is_multiple=True,
)
def slowtest3(self):
print "Generate a model for M31, using defaults (100k disk, 50k bulge, 200k halo) - Nbody units"
halo_number_of_particles = 20000
particles = new_galactics_model(halo_number_of_particles, do_scale = True,
bulge_number_of_particles=5000, disk_number_of_particles=10000)
self.assertEquals(len(particles), 35000)
self.assertAlmostEquals(particles.total_mass(), 1.0 | nbody_system.mass)
self.assertAlmostEquals(particles.kinetic_energy(), 0.25 | nbody_system.energy)
write_set_to_file(particles, os.path.join(get_path_to_results(), 'M31_galactICs.amuse'), 'amuse')
def test26(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test26"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
p = Grid(20)
io.write_set_to_file(p, output_file, format='amuse', version = self.store_version())
os.remove(output_file)
def test23(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test23"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
particles = Particles()
io.write_set_to_file(particles, output_file, format='amuse', version = self.store_version())
particles_from_file = io.read_set_from_file(output_file, format='amuse', version = self.store_version())
self.assertEquals(len(particles_from_file), 0)
def relax_in_isolation(giant_in_sph, sph_code, output_base_name):
total_mass = giant_in_sph.gas_particles.total_mass() + giant_in_sph.core_particle.mass
total_radius = max(giant_in_sph.gas_particles.position.lengths_squared()).sqrt()
dynamical_timescale = (total_radius**3 / (2 * constants.G * total_mass)).sqrt().as_quantity_in(units.day)
t_end = (10.0 * dynamical_timescale).as_quantity_in(units.day)
n_steps = 100
hydro_code_options = dict(number_of_workers=2, redirection='file', redirect_file='hydrodynamics_code_relax_out.log')
unit_converter = ConvertBetweenGenericAndSiUnits(total_radius, total_mass, t_end)
hydrodynamics = sph_code(unit_converter, **hydro_code_options)
hydrodynamics.parameters.epsilon_squared = giant_in_sph.core_radius**2
hydrodynamics.parameters.max_size_timestep = t_end
hydrodynamics.parameters.time_max = 1.1 * t_end
hydrodynamics.parameters.time_limit_cpu = 7.0 | units.day
hydrodynamics.gas_particles.add_particles(giant_in_sph.gas_particles)
hydrodynamics.dm_particles.add_particle(giant_in_sph.core_particle)
potential_energies = hydrodynamics.potential_energy.as_vector_with_length(1).as_quantity_in(units.erg)
kinetic_energies = hydrodynamics.kinetic_energy.as_vector_with_length(1).as_quantity_in(units.erg)
thermal_energies = hydrodynamics.thermal_energy.as_vector_with_length(1).as_quantity_in(units.erg)
print " Relaxing for", t_end, "(10 * dynamical timescale)"
times = (t_end * range(1, n_steps+1) / n_steps).as_quantity_in(units.day)
for i_step, time in enumerate(times):
hydrodynamics.evolve_model(time)
print " Relaxed for:", time
potential_energies.append(hydrodynamics.potential_energy)
kinetic_energies.append(hydrodynamics.kinetic_energy)
thermal_energies.append(hydrodynamics.thermal_energy)
hydrodynamics.gas_particles.copy_values_of_attributes_to(
['mass', 'x','y','z', 'vx','vy','vz', 'u', 'h_smooth'],
giant_in_sph.gas_particles)
giant_in_sph.core_particle.position = hydrodynamics.dm_particles[0].position
giant_in_sph.core_particle.velocity = hydrodynamics.dm_particles[0].velocity
hydrodynamics.stop()
snapshotfile = output_base_name + "_gas.amuse"
write_set_to_file(giant_in_sph.gas_particles, snapshotfile, format='amuse')
shutil.copy(snapshotfile, os.path.join("..", "giant_models"))
snapshotfile = output_base_name + "_core.amuse"
# temporarily store core_radius on the core particle
giant_in_sph.core_particle.radius = giant_in_sph.core_radius
write_set_to_file(giant_in_sph.core_particle.as_set(), snapshotfile, format='amuse')
giant_in_sph.core_particle.radius = 0.0 | units.m
shutil.copy(snapshotfile, os.path.join("..", "giant_models"))
energy_evolution_plot(times, kinetic_energies, potential_energies, thermal_energies,
figname = output_base_name + "_energy_evolution.png")
def test10(self):
print "Test 10: User interface (write_set_to_file and read_set_from_file)"
particles = datamodel.Particles(2)
particles.a = [1, 4] | units.none
particles.b = [2, 5] | units.m
particles.c = [3, 6] | units.kg / units.m**3
io.write_set_to_file(particles, "test_textio.csv","csv")
read_particles = io.read_set_from_file("test_textio.csv", format = "csv")
self.assertEquals(len(read_particles), 2)
self.assertEquals(read_particles.a, [1, 4] | units.none)
self.assertEquals(read_particles.b, [2, 5] | units.m)
self.assertEquals(read_particles.c, [3, 6] | units.kg / units.m**3)
os.remove("test_textio.csv")
def test29(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test29"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
particles = Grid(10)
particles.attribute1 = "normal"
particles.attribute2 = u"unicode"
io.write_set_to_file(particles,output_file, format='amuse', version = self.store_version())
output = io.read_set_from_file(output_file, format='amuse')
self.assertEquals(output[0].attribute2, u"unicode")
def test1(self):
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | nbody_system.mass
x.radius = [3.0, 4.0] | nbody_system.length
x.position = [[1,2,3], [3,5,6]] | nbody_system.length
x.velocity = [[1,2,3], [3,5,6]] | nbody_system.speed
io.write_set_to_file(x, "test.tsf","tsf")
y = io.read_set_from_file("test.tsf","tsf")
self.assertAlmostEquals(x.mass, y.mass, 8)
# self.assertAlmostEquals(x.radius, y.radius, 8)
self.assertAlmostEquals(x.position, y.position,8)
self.assertAlmostEquals(x.velocity, y.velocity,8)
os.remove("test.tsf")
def test59(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test59"+self.store_version()+".h5")
if os.path.exists(output_file):
os.remove(output_file)
g=Grid(10,10)
p=Particles(2)
p[0].sub=g[0:5]
io.write_set_to_file(p, output_file,"amuse", version=self.store_version())
z = io.read_set_from_file(output_file,"amuse")
os.remove(output_file)
def hybrid_self_gravitating_cluster(Ncl, Mstar, Rcl, t_end, dt,
Mcut, filename):
stars_below_cut, stars_above_cut, nbody_converter =\
create_cluster(Ncl, Mstar, Rcl, Mcut)
gravity_low_mass, gravity_high_mass, bridge, channels = \
setup_codes(stars_below_cut, stars_above_cut, nbody_converter)
time = 0. | units.Myr
write_set_to_file((stars_below_cut.savepoint(time), stars_above_cut.savepoint(time)), filename,
'amuse', names=("stars_below_cut", "stars_above_cut"), append_to_file=False)
# Calculate initial energy
Etot_init_low_mass = gravity_low_mass.kinetic_energy +\
gravity_low_mass.potential_energy
Etot_init_high_mass = gravity_high_mass.kinetic_energy +\
gravity_high_mass.potential_energy
Etot_init = Etot_init_low_mass + Etot_init_high_mass
print "Saving data to", filename
while time < (t_end - dt/2.):
Etot_prev = Etot_init
time += dt
print "Evolving to", time
bridge.evolve_model(time)
channels.copy()
# Calculate relative energy error
Ekin_low_mass = gravity_low_mass.kinetic_energy
Epot_low_mass = gravity_low_mass.potential_energy
Etot_low_mass = Ekin_low_mass + Epot_low_mass
Ekin_high_mass = gravity_high_mass.kinetic_energy
Epot_high_mass = gravity_high_mass.potential_energy
Etot_high_mass = Ekin_high_mass + Epot_high_mass
Etot = Etot_low_mass + Etot_high_mass
print "dE =", (Etot_init - Etot) / Etot, "ddE =",(Etot_prev - Etot) / Etot
dE = (Etot_init - Etot) / Etot
if isnan(dE): # numpy function
dE = 0
# Save data
write_set_to_file((stars_below_cut.savepoint(time), stars_above_cut.savepoint(time)),
filename, 'amuse', names=("stars_below_cut", "stars_above_cut"))
def test2(self):
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | nbody_system.mass
x.radius = [3.0, 4.0] | nbody_system.length
x.position = [[1,2,3], [3,5,6]] | nbody_system.length
x.velocity = [[1,2,3], [3,5,6]] | nbody_system.speed
io.write_set_to_file(x, "test.dyn","dyn")
y = io.read_set_from_file("test.dyn","dyn")
self.assertAlmostEquals(x.mass, y.mass, 8)
self.assertAlmostEquals(x.position, y.position,8)
self.assertAlmostEquals(x.velocity, y.velocity,8)
self.assertRaises(AttributeError, lambda: y.radius,
expected_message = "You tried to access attribute 'radius' but this "
"attribute is not defined for this set.")
os.remove("test.dyn")
def test57(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test26"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
x = Particles(2)
x.sub = None
x[0].sub = Particles(2)
io.write_set_to_file(x, output_file,"amuse", version=self.store_version())
y = io.read_set_from_file(output_file,"amuse")
self.assertEqual(len(x), len(y))
self.assertEqual(len(x[0].sub), len(y[0].sub))
self.assertTrue(y[1].sub==x[1].sub==None)
os.remove(output_file)
def test10(self):
stellar_evolution = SSE()
stellar_evolution.commit_parameters()
stars = Particles(10)
stars.mass = 1.0 | units.MSun
stellar_evolution.particles.add_particles(stars)
self.assertEquals(stellar_evolution.particles._factory_for_new_collection(), Particles)
filename = os.path.join(get_path_to_results(), "test.h5")
if os.path.exists(filename):
os.remove(filename)
io.write_set_to_file(stellar_evolution.particles, filename, 'hdf5')
stored_stars = io.read_set_from_file(filename, 'hdf5')
self.assertEquals(len(stars), len(stored_stars))
self.assertAlmostRelativeEquals(stars.mass, stored_stars.mass)
def test29(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test29" + self.store_version() + ".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
particles = Grid(10)
particles.attribute1 = "normal"
particles.attribute2 = "unicode"
io.write_set_to_file(particles,
output_file,
format='amuse',
version=self.store_version())
output = io.read_set_from_file(output_file, format='amuse')
self.assertEqual(output[0].attribute2, "unicode")
def test5(self):
print("Testing HDF5 io")
if os.path.exists("test_unit.hdf5"):
os.remove("test_unit.hdf5")
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | units.kg
x.radius = [3.0, 4.0] | units.m
x.position = [[1, 2, 3], [3, 5, 6]] | units.m
x.velocity = [[1, 2, 3], [3, 5, 6]] | units.m / units.s
io.write_set_to_file(x, "test_unit.hdf5", "hdf5")
y = io.read_set_from_file("test_unit.hdf5", "hdf5")
self.assertAlmostEqual(x.mass, y.mass, 8)
self.assertAlmostEqual(x.radius, y.radius, 8)
self.assertAlmostEqual(x.position, y.position, 8)
self.assertAlmostEqual(x.velocity, y.velocity, 8)
os.remove("test_unit.hdf5")
def test23(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test23" + self.store_version() + ".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
particles = Particles()
io.write_set_to_file(particles,
output_file,
format='amuse',
version=self.store_version())
particles_from_file = io.read_set_from_file(
output_file, format='amuse', version=self.store_version())
self.assertEqual(len(particles_from_file), 0)
def slowtest3(self):
print "Generate a model for M31, using defaults (100k disk, 50k bulge, 200k halo) - Nbody units"
halo_number_of_particles = 20000
particles = new_galactics_model(halo_number_of_particles,
do_scale=True,
bulge_number_of_particles=5000,
disk_number_of_particles=10000)
self.assertEquals(len(particles), 35000)
self.assertAlmostEquals(particles.total_mass(),
1.0 | nbody_system.mass)
self.assertAlmostEquals(particles.kinetic_energy(),
0.25 | nbody_system.energy)
write_set_to_file(
particles,
os.path.join(get_path_to_results(), 'M31_galactICs.amuse'),
'amuse')
def test11(self):
particles = datamodel.Particles(2)
particles.a = [1, 4]
particles.b = [2, 5] | units.m
particles.c = [3, 6] | units.kg / units.m**3
io.write_set_to_file(particles,
"test_textio.csv",
"csv",
attribute_type=(None, units.kg / units.m**3,
units.m),
attribute_name=('a', 'c', 'b'))
read_particles = io.read_set_from_file("test_textio.csv", format="csv")
self.assertEquals(len(read_particles), 2)
self.assertEquals(read_particles.a, [1, 4])
self.assertEquals(read_particles.b, [2, 5] | units.m)
self.assertEquals(read_particles.c, [3, 6] | units.kg / units.m**3)
os.remove("test_textio.csv")
def test3(self):
x = datamodel.Particles(2)
convert = nbody_system.nbody_to_si(1 | units.kg, 2 | units.m)
x.mass = [1.0, 2.0] | units.kg
x.radius = [3.0, 4.0] | units.m
x.position = [[1,2,3], [3,5,6]] | units.m
x.velocity = [[1,2,3], [3,5,6]] | units.m / units.s
io.write_set_to_file(x, "test_unit.dyn","dyn", nbody_to_si_converter = convert)
y = io.read_set_from_file("test_unit.dyn","dyn", nbody_to_si_converter = convert)
self.assertAlmostEqual(x.mass, y.mass, 8)
self.assertAlmostEqual(x.position, y.position,8)
self.assertAlmostEqual(x.velocity, y.velocity,8)
self.assertRaises(AttributeError, lambda: y.radius,
expected_message = "You tried to access attribute 'radius' but this "
"attribute is not defined for this set.")
os.remove("test_unit.dyn")
def radhydro_evolve(sph, rad, tend, dt):
"""
evolve function to co-evolve sph under radiative feedback from rad
the evolve proceeds as follows (a form of leapfrog integrator):
- 1/2 step sph
- update positions and densities in rad
- full step rad transfer
- update internal energies in sph
- 1/2 step sph
tend=end time
dt=time step
this function dump snapshots in file dump-..
"""
i = 0
write_set_to_file(sph.gas_particles,
"dump-%6.6i" % i,
"amuse",
append_to_file=False)
t = 0. | units.Myr
while t < tend - dt / 2:
print t
print "sph1"
sph.evolve_model(t + dt / 2)
print "rad"
update_pos_rho(rad, sph.gas_particles)
rad.evolve_model(t + dt)
update_u(sph, rad.particles)
print "sph2"
sph.evolve_model(t + dt)
t += dt
i += 1
write_set_to_file(sph.gas_particles,
"dump-%6.6i" % i,
"amuse",
append_to_file=False)
sph.stop()
rad.stop()
def test20(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test20"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
shape = 10, 10, 10
p = Grid(*shape)
p.mass = (numpy.asarray([x * 2.0 for x in range(p.size)])).reshape(shape)
io.write_set_to_file(p, output_file, "hdf5", version = self.store_version())
loaded = io.read_set_from_file(output_file, "hdf5", version = self.store_version())
loaded = self.get_version_in_store(loaded)
self.assertAlmostRelativeEquals(p.mass[0][1][2], 24)
self.assertAlmostRelativeEquals(p[0][1][2].mass, 24)
def generate_initial_conditions(
number_of_stars = 10000,
number_of_gas_particles = 2*10**6,
star_formation_efficiency = 0.1,
virial_radius = 0.33 | units.parsec,
virial_ratio = 1.0,
use_fractal = False):
numpy.random.seed(12345678)
seed_fractal = 312357271
masses = new_kroupa_mass_distribution(number_of_stars)
total_stellar_mass = masses.sum()
total_mass = total_stellar_mass / star_formation_efficiency
converter = nbody_system.nbody_to_si(total_mass, virial_radius)
if use_fractal:
stars = new_fractal_cluster_model(number_of_stars, convert_nbody=converter, do_scale=False, fractal_dimension=1.6, random_seed=seed_fractal)
else:
stars = new_plummer_model(number_of_stars, convert_nbody=converter, do_scale=False)
stars.mass = masses
stars.move_to_center()
print "scaling positions to match virial_radius"
stars.position *= virial_radius / stars.virial_radius()
print "scaling velocities to match virial_ratio"
stars.velocity *= numpy.sqrt(virial_ratio * converter.to_si(0.5|nbody_system.energy) * star_formation_efficiency / stars.kinetic_energy())
print "new_gas_plummer_distribution"
gas = new_gas_plummer_distribution(
number_of_gas_particles,
total_mass = (total_mass - total_stellar_mass),
virial_radius = virial_radius,
type = "fcc")
gas.h_smooth = 0.0 | units.parsec
filename = "YSC_{0}_stars{1}_gas{2}k_".format("fractal" if use_fractal else "plummer",
number_of_stars, number_of_gas_particles/1000)
print "Writing initial conditions to", filename, "+ stars/gas.amuse"
write_set_to_file(stars, filename+"stars.amuse", "amuse", append_to_file=False)
write_set_to_file(gas, filename+"gas.amuse", "amuse", append_to_file=False)
with open(filename+"info.pkl", "wb") as outfile:
cPickle.dump([converter], outfile)
return stars, gas, filename
def test14(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test14"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
stars = Particles(2)
stars.x = 1.0 | units.km
stars.md = [[[1,3],[2,4],[3,5]],[[4,6],[5,7],[6,8]]]
io.write_set_to_file(stars, output_file, "hdf5", version = self.store_version())
loaded = io.read_set_from_file(output_file, "hdf5", version = self.store_version())
self.assertEquals(loaded[0].md, [[1,3],[2,4],[3,5]])
self.assertEquals(loaded[1].md, [[4,6],[5,7],[6,8]])
self.assertEquals(loaded.md[0], [[1,3],[2,4],[3,5]])
self.assertEquals(loaded.md[1], [[4,6],[5,7],[6,8]])
def test10(self):
stellar_evolution = SSE()
stellar_evolution.commit_parameters()
stars = Particles(10)
stars.mass = 1.0 | units.MSun
stellar_evolution.particles.add_particles(stars)
self.assertEqual(
stellar_evolution.particles._factory_for_new_collection(),
Particles)
filename = os.path.join(get_path_to_results(), "test.h5")
if os.path.exists(filename):
os.remove(filename)
io.write_set_to_file(stellar_evolution.particles, filename, 'hdf5')
stored_stars = io.read_set_from_file(filename, 'hdf5')
self.assertEqual(len(stars), len(stored_stars))
self.assertAlmostRelativeEquals(stars.mass, stored_stars.mass)
def test30(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test30"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
particles = Particles(10)
particles.mass = 10 | units.kg
io.write_set_to_file(particles,output_file, format='amuse', version = self.store_version())
output = io.read_set_from_file(output_file, format='amuse', allow_writing = True)
output.new_attribute= 2 * output.mass
self.assertEquals(output.new_attribute, 2 * output.mass)
output = output.iter_history().next()
output.new_attribute= 2 * output.mass
self.assertEquals(output.new_attribute, 2 * output.mass)
def test2(self):
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | nbody_system.mass
x.radius = [3.0, 4.0] | nbody_system.length
x.position = [[1, 2, 3], [3, 5, 6]] | nbody_system.length
x.velocity = [[1, 2, 3], [3, 5, 6]] | nbody_system.speed
io.write_set_to_file(x, "test.dyn", "dyn")
y = io.read_set_from_file("test.dyn", "dyn")
self.assertAlmostEquals(x.mass, y.mass, 8)
self.assertAlmostEquals(x.position, y.position, 8)
self.assertAlmostEquals(x.velocity, y.velocity, 8)
self.assertRaises(
AttributeError,
lambda: y.radius,
expected_message="You tried to access attribute 'radius' but this "
"attribute is not defined for this set.")
os.remove("test.dyn")
def evolve_system(system, t_end, n_steps):
times = (t_end * range(1, n_steps+1) / n_steps).as_quantity_in(units.day)
for i_step, time in enumerate(times):
system.evolve_model(time)
print " Evolved to:", time,
figname = os.path.join("plots", "hydro_giant_{0:=04}.png".format(i_step))
print " - Hydroplot saved to: ", figname
pynbody_column_density_plot(system.gas_particles, width=30|units.AU, vmin=26, vmax=33)
scatter(system.dm_particles.x, system.dm_particles.y, c="w")
pyplot.savefig(figname)
pyplot.close()
file_base_name = os.path.join("snapshots", "hydro_giant_{0:=04}_".format(i_step))
write_set_to_file(system.gas_particles, file_base_name+"gas.amuse", format='amuse')
write_set_to_file(system.dm_particles, file_base_name+"dm.amuse", format='amuse')
system.stop()
def test60(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test60"+self.store_version()+".h5")
if os.path.exists(output_file):
os.remove(output_file)
x = Particles(2)
x.mass = 1 | units.kg
y = Particles(2)
y.mass = 2 | units.kg
io.write_set_to_file((x,y), output_file,"amuse", version=self.store_version(), names = ("x", "y"))
z = io.read_set_from_file(output_file,"amuse")
os.remove(output_file)
self.assertTrue("x" in z)
self.assertTrue("y" in z)
self.assertAlmostRelativeEquals(z["x"].mass, 1 | units.kg)
self.assertAlmostRelativeEquals(z["y"].mass, 2 | units.kg)
def test61(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test61"+self.store_version()+".h5")
if os.path.exists(output_file):
os.remove(output_file)
x = Particles(10)
x.mass = 1 | units.kg
y = Particles(20)
y.id = range(20)
x[0].y = y[5:7]
io.write_set_to_file(x, output_file,"amuse", version=self.store_version())
z = io.read_set_from_file(output_file,"amuse")
self.assertEquals(x[0].y[0].id, 5)
self.assertEquals(z[0].y[0].id, 5)
os.remove(output_file)
def test13(self):
p = datamodel.Particles(100)
p.a = numpy.arange(0, 1, 0.01) | units.m
path = os.path.abspath(
os.path.join(self.get_path_to_results(), "test.csv"))
io.write_set_to_file(p,
path,
"amuse-txt",
attribute_names=('a'),
maximum_number_of_lines_buffered=10,
key_in_column=0)
p2 = io.read_set_from_file(path,
"txt",
maximum_number_of_lines_buffered=10,
key_in_column=0)
self.assertEquals(p2.key, p.key)
self.assertAlmostRelativeEquals(p2.a, p.a)
def test12(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test12"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
shape = 10, 10, 10
p = Grid(*shape)
p.mass = ([x * 2.0 for x in range(p.size)] | units.kg).reshape(shape)
p.model_time = 2.0 | units.s
io.write_set_to_file(p.savepoint(timestamp = 2 | units.Myr, scale = 1 | units.kg), output_file, "hdf5", version = self.store_version())
loaded_particles = io.read_set_from_file(output_file, "hdf5", version = self.store_version())
loaded_particles = self.get_version_in_store(loaded_particles)
a = loaded_particles.collection_attributes
self.assertAlmostRelativeEquals(a.timestamp, 2 | units.Myr)
self.assertAlmostRelativeEquals(a.scale, 1 | units.kg)
def slowtest4(self):
print "Generate a model for a disk galaxy (10k disk, no bulge, 20k halo) - SI units"
halo_number_of_particles = 20000
converter = generic_unit_converter.ConvertBetweenGenericAndSiUnits(constants.G, 1.0e12 | units.MSun, 50.0 | units.kpc)
particles = new_galactics_model(halo_number_of_particles, disk_number_of_particles=10000,
generate_bulge_flag=False, do_scale=True, unit_system_converter=converter)
self.assertEquals(len(particles), 30000)
self.assertAlmostRelativeEquals(particles.total_mass(), 1.0e12 | units.MSun, 10)
self.assertAlmostRelativeEquals(particles.kinetic_energy(), converter.to_si(0.25 | nbody_system.energy), 10)
disk = particles[:10000]
self.assertAlmostRelativeEquals(disk.total_mass(), 2.156e10 | units.MSun, 3)
self.assertAlmostRelativeEquals(disk.position.lengths_squared().amax().sqrt().in_(units.kpc),
15.584 | units.kpc, 3)
self.assertAlmostRelativeEquals(disk.position.std(axis=0).in_(units.kpc), [3.5934, 3.6768, 0.17078] | units.kpc, 3)
write_set_to_file(particles, os.path.join(get_path_to_results(), 'disk_galactICs.amuse'), 'amuse')
def test11(self):
p = datamodel.Particles(200)
p.a = 2 | units.m
io.write_set_to_file(p,
"test.csv",
"txt",
attribute_names=('a'),
attribute_types=(units.m, ),
maximum_number_of_lines_buffered=10,
key_in_column=0)
p2 = io.read_set_from_file("test.csv",
"txt",
attribute_names=('a'),
attribute_types=(units.m, ),
maximum_number_of_lines_buffered=10,
key_in_column=0)
self.assertEquals(p2.key, p.key)
self.assertAlmostRelativeEquals(p2.a, p.a)
def test11(self):
print "Testing saving/loading timestamp in NEMO"
x = datamodel.Particles(2)
convert = nbody_system.nbody_to_si(1 | units.kg, 2 | units.m)
x.mass = [1.0, 2.0] | units.kg
x.position = [[1, 2, 3], [3, 5, 6]] | units.m
x.velocity = [[1, 2, 3], [3, 5, 6]] | units.m / units.s
current_time = 1.0 | units.Myr
io.write_set_to_file(x.savepoint(current_time),
"time_test_unit.tsf",
"tsf",
nbody_to_si_converter=convert)
y = io.read_set_from_file("time_test_unit.tsf",
"tsf",
nbody_to_si_converter=convert)
self.assertAlmostEquals(current_time,
y.previous_state().get_timestamp(),
8,
in_units=units.Myr)
self.assertAlmostEquals(x.mass, y.mass, 8)
self.assertAlmostEquals(x.position, y.position, 8)
self.assertAlmostEquals(x.velocity, y.velocity, 8)
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | nbody_system.mass
x.position = [[1, 2, 3], [3, 5, 6]] | nbody_system.length
x.velocity = [[1, 2, 3], [3, 5, 6]] | nbody_system.speed
current_time = 1.0 | nbody_system.time
io.write_set_to_file(x.savepoint(current_time), "time_test_unit.tsf",
"tsf")
y = io.read_set_from_file("time_test_unit.tsf", "tsf")
self.assertAlmostEquals(current_time,
y.previous_state().get_timestamp(),
8,
in_units=nbody_system.time)
self.assertAlmostEquals(x.mass, y.mass, 8)
self.assertAlmostEquals(x.position, y.position, 8)
self.assertAlmostEquals(x.velocity, y.velocity, 8)
os.remove("time_test_unit.tsf")
def evolve_system(coupled_system, t_end, n_steps):
times = (t_end * range(1, n_steps + 1) / n_steps).as_quantity_in(units.day)
hydro = coupled_system.codes[
0].code # only calculate potential energy for the giant (SPH particles)
potential_energies = hydro.potential_energy.as_vector_with_length(
1).as_quantity_in(units.J)
kinetic_energies = hydro.kinetic_energy.as_vector_with_length(
1).as_quantity_in(units.J)
thermal_energies = hydro.thermal_energy.as_vector_with_length(
1).as_quantity_in(units.J)
giant = coupled_system.particles
giant_initial_position = giant.center_of_mass()
giant_initial_velocity = giant.center_of_mass_velocity()
for i_step, time in enumerate(times):
giant.position += giant_initial_position - giant.center_of_mass()
giant.velocity = (giant.velocity - giant.center_of_mass_velocity()) * (
i_step * 1.0 / n_steps) + giant_initial_velocity
print " Evolving...",
coupled_system.evolve_model(time)
print " Evolved to:", time,
potential_energies.append(hydro.potential_energy)
kinetic_energies.append(hydro.kinetic_energy)
thermal_energies.append(hydro.thermal_energy)
print " Energies calculated"
density_plot(coupled_system, i_step)
snapshotfile = "hydro_giant_gas.amuse"
write_set_to_file(coupled_system.gas_particles,
snapshotfile,
format='amuse')
snapshotfile = "hydro_giant_dm.amuse"
write_set_to_file(coupled_system.dm_particles,
snapshotfile,
format='amuse')
coupled_system.stop()
energy_evolution_plot(times[:len(kinetic_energies) - 1], kinetic_energies,
potential_energies, thermal_energies)
def test16(self):
convert_nbody = nbody_system.nbody_to_si(1.0 | units.MSun,
149.5e6 | units.km)
smalln = SmallN(convert_nbody)
smalln.initialize_code()
smalln.dt_dia = 5000
stars = self.new_system_of_sun_and_earth()
moon = datamodel.Particle()
moon.mass = units.kg(7.3477e22)
moon.radius = units.km(1737.10)
moon.position = units.km(numpy.array((149.5e6 + 384.399, 0.0, 0.0)))
moon.velocity = units.ms(numpy.array((0.0, 29800 + 1022, 0.0)))
stars.add_particle(moon)
earth = stars[1]
smalln.particles.add_particles(stars)
smalln.evolve_model(365.0 | units.day)
smalln.update_particle_tree()
smalln.update_particle_set()
self.assertEqual(len(smalln.particles), 5)
self.assertEarthAndMoonWasDetectedAsBinary(smalln.particles, stars)
inmemory = smalln.particles.copy()
self.assertEarthAndMoonWasDetectedAsBinary(inmemory, stars)
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "newsmalln-test16.hdf5")
if os.path.exists(output_file):
os.remove(output_file)
io.write_set_to_file(smalln.particles, output_file, "hdf5")
fromfile = io.read_set_from_file(output_file, "hdf5")
self.assertEarthAndMoonWasDetectedAsBinary(fromfile, stars)
smalln.stop()
def test19(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test19" + self.store_version() + ".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
stars = Particles(2)
stars[0].x = 1.0 | units.km
stars[1].x = 2.0 | units.km
binaries = Particles(1)
binaries[0].y = 1.0 | units.km
binaries[0].child1 = stars[0]
binaries[0].child2 = stars[1]
stars[0].parent = binaries[0]
stars[1].parent = binaries[0]
self.assertEqual(binaries[0].child1, stars[0])
self.assertEqual(binaries[0].child2, stars[1])
self.assertEqual(binaries[0].child1.parent, binaries[0])
self.assertEqual(binaries[0].child2.parent, binaries[0])
io.write_set_to_file([binaries, stars],
output_file,
"hdf5",
names=['binaries', 'children'],
version=self.store_version())
loader_binaries, loaded_stars = io.read_set_from_file(
output_file,
"hdf5",
names=['binaries', 'children'],
version=self.store_version())
self.assertEqual(loader_binaries[0].child1.key, stars[0].key)
self.assertEqual(loader_binaries[0].child2.key, stars[1].key)
self.assertEqual(loader_binaries[0].child1, loaded_stars[0])
self.assertEqual(loader_binaries[0].child2, loaded_stars[1])
self.assertEqual(loader_binaries[0].child1.parent, loader_binaries[0])
self.assertEqual(loader_binaries[0].child2.parent, loader_binaries[0])
def test59(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test59" + self.store_version() + ".h5")
if os.path.exists(output_file):
os.remove(output_file)
g = Grid(10, 10)
p = Particles(2)
p[0].sub = g[0:5]
io.write_set_to_file(p,
output_file,
"amuse",
version=self.store_version())
z = io.read_set_from_file(output_file, "amuse")
os.remove(output_file)
def test22(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test22"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
stars = Particles(2)
stars.x = 1.0 | units.km
overlay = ParticlesOverlay(stars)
overlay.y = 2.0 | units.km
io.write_set_to_file(overlay, output_file, "hdf5", version = self.store_version())
loaded = io.read_set_from_file(output_file, "hdf5", close_file = True, version = self.store_version())
self.assertEquals(loaded[0].x, 1.0 | units.km)
self.assertEquals(loaded[1].y, 2.0 | units.km)
def test7(self):
print("Testing HDF5 io with a ParticlesSuperset")
if os.path.exists("test_unit.hdf5"):
os.remove("test_unit.hdf5")
set1 = datamodel.Particles(2)
set2 = datamodel.Particles(2)
superset = datamodel.ParticlesSuperset([set1, set2])
superset.mass = [1.0, 2.0, 3.0, 4.0] | units.kg
superset.radius = [3.0, 4.0, 5.0, 6.0] | units.m
superset.position = [[1,2,3], [3,5,6], [3,2,1], [-3,-5,-6]] | units.m
superset.velocity = [[1,2,3], [3,5,6], [3,2,1], [-3,-5,-6]] | units.m / units.s
io.write_set_to_file(superset, "test_unit.hdf5","hdf5")
y = io.read_set_from_file("test_unit.hdf5","hdf5")
self.assertAlmostEqual(superset.mass, y.mass, 8)
self.assertAlmostEqual(superset.radius, y.radius, 8)
self.assertAlmostEqual(superset.position, y.position,8)
self.assertAlmostEqual(superset.velocity, y.velocity,8)
os.remove("test_unit.hdf5")
def test28(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test28" + self.store_version() + ".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
stars = Particles(2)
stars[0].x = 1.0 | units.km
stars[1].x = 2.0 | units.km
binaries = Particles(1)
binaries[0].y = 1.0 | units.km
binaries[0].child1 = stars[0]
binaries[0].child2 = stars[1]
stars[0].parent = binaries[0]
stars[1].parent = binaries[0]
self.assertEqual(binaries[0].child1, stars[0])
self.assertEqual(binaries[0].child2, stars[1])
self.assertEqual(binaries[0].child1.parent, binaries[0])
self.assertEqual(binaries[0].child2.parent, binaries[0])
io.write_set_to_file(binaries,
output_file,
"hdf5",
version=self.store_version())
with io.read_set_from_file(output_file,
"hdf5",
version=self.store_version(),
return_context=True) as loaded:
self.assertEqual(loaded[0].child1.key, stars[0].key)
self.assertEqual(loaded[0].child2.key, stars[1].key)
self.assertEqual(loaded[0].child1, stars[0])
self.assertEqual(loaded[0].child2, stars[1])
self.assertEqual(loaded[0].child1.parent, loaded[0])
self.assertEqual(loaded[0].child2.parent, loaded[0])
def test57(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test26" + self.store_version() + ".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
x = Particles(2)
x.sub = None
x[0].sub = Particles(2)
io.write_set_to_file(x,
output_file,
"amuse",
version=self.store_version())
y = io.read_set_from_file(output_file, "amuse")
self.assertEqual(len(x), len(y))
self.assertEqual(len(x[0].sub), len(y[0].sub))
self.assertTrue(y[1].sub == x[1].sub == None)
os.remove(output_file)
def test53(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path, "test53"+self.store_version()+".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
stars = Particles(2)
stars[0].x = 1.0 | units.km
stars[1].x = 2.0 | units.km
gas = Grid(2,3)
gas.y = [[1.0, 2.0, 3.0], [4.0,5.0,6.0]] | units.km
stars[0].gas = gas
io.write_set_to_file(stars, output_file, "hdf5", version = self.store_version())
loaded_stars = io.read_set_from_file(output_file, "hdf5", version = self.store_version())
self.assertAlmostRelativeEquals(loaded_stars[0].gas[0][0].y,1.0 | units.km)
self.assertAlmostRelativeEquals(loaded_stars[0].gas[0][2].y,3.0 | units.km)
self.assertEquals(loaded_stars[1].gas, None)
