def test13(self):
print("Test convert_gadget_w_to_velocity and return_header for Gadget read_set_from_file")
directory_name = os.path.dirname(__file__)
filename = os.path.join(directory_name, 'tiny_lcdm_data_littleendian.dat')
data = io.read_set_from_file(filename, format='gadget', return_header=False, convert_gadget_w_to_velocity=False) # (default)
self.assertTrue(isinstance(data, tuple))
self.assertEqual(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry)")
self.assertEqual(len(data.gas), 32)
self.assertEqual(len(data.halo), 32)
self.assertEqual(data.gas[0].key, 1)
self.assertEqual(data.halo[0].key, 32**3 + 1)
self.assertAlmostRelativeEquals(data.gas[:3].position, [[395.23443604, 395.75210571, 1244.31152344],
[310.17266846, 440.21728516, 2817.06396484], [191.95669556, 465.57223511, 4430.20068359]] | nbody_system.length, 7)
data_converted = io.read_set_from_file(filename, format='gadget', return_header=True, convert_gadget_w_to_velocity=True)
self.assertTrue(isinstance(data_converted, tuple))
self.assertEqual(data_converted.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry, "
"Npart, Massarr, Time, Redshift, FlagSfr, FlagFeedback, Nall, FlagCooling, "
"NumFiles, BoxSize, Omega0, OmegaLambda, HubbleParam, FlagAge, FlagMetals, "
"NallHW, flag_entr_ics)")
self.assertEqual(len(data_converted.gas), 32)
self.assertEqual(len(data_converted.halo), 32)
self.assertEqual(data_converted.gas[0].key, 1)
self.assertEqual(data_converted.halo[0].key, 32**3 + 1)
self.assertEqual(data_converted.Npart, (32, 32, 0, 0, 0, 0))
self.assertEqual(data_converted.Time, 1/11.0)
self.assertEqual(data_converted.Redshift, 10.0)
self.assertEqual(data.gas.position, data_converted.gas.position)
self.assertAlmostRelativeEquals(data.gas.velocity, math.sqrt(data_converted.Time) * data_converted.gas.velocity, 7)
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.assertEqual(len(gas), 1472)
self.assertEqual(len(halo), 0)
self.assertEqual(gas[0].key,1)
self.assertEqual(gas[1].key,2)
self.assertEqual(gas[2].key,3)
self.assertEqual(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.assertEqual(len(gas), 1472)
self.assertEqual(len(halo), 0)
self.assertEqual(gas[0].key,1)
self.assertEqual(gas[1].key,2)
self.assertEqual(gas[2].key,3)
self.assertEqual(gas[1471].key,1472)
finally:
if os.path.exists(outputfilename):
os.remove(outputfilename)
def test12(self):
print("Test return_header for Gadget read_set_from_file")
directory_name = os.path.dirname(__file__)
filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
data = io.read_set_from_file(filename, format='gadget', return_header=False) # (default)
self.assertTrue(isinstance(data, tuple))
self.assertEqual(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry)")
data = io.read_set_from_file(filename, format='gadget', return_header=True)
self.assertTrue(isinstance(data, tuple))
self.assertEqual(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry, "
"Npart, Massarr, Time, Redshift, FlagSfr, FlagFeedback, Nall, FlagCooling, "
"NumFiles, BoxSize, Omega0, OmegaLambda, HubbleParam, FlagAge, FlagMetals, "
"NallHW, flag_entr_ics)")
self.assertEqual(len(data.gas), 1472)
self.assertEqual(len(data.halo), 0)
self.assertEqual(data.gas[0].key,1)
self.assertEqual(data.gas[1].key,2)
self.assertEqual(data.gas[2].key,3)
self.assertEqual(data.gas[1471].key,1472)
self.assertAlmostRelativeEquals(data.gas[0:5].x,[-0.0713372901082, 0.0713372901082, -0.21178227663, -0.0698266476393, 0.0698266476393] | nbody_system.length, 7)
self.assertAlmostRelativeEquals(data.gas[0:5].u, [0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451] | (nbody_system.length / nbody_system.time)**2, 7 )
self.assertEqual(data.Npart, (1472, 0, 0, 0, 0, 0))
self.assertEqual(data.Time, 0.0)
self.assertEqual(data.Redshift, 0.0)
def test12(self):
print "Test return_header for Gadget read_set_from_file"
directory_name = os.path.dirname(__file__)
filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
data = io.read_set_from_file(filename, format='gadget', return_header=False) # (default)
self.assertTrue(isinstance(data, tuple))
self.assertEquals(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry)")
data = io.read_set_from_file(filename, format='gadget', return_header=True)
self.assertTrue(isinstance(data, tuple))
self.assertEquals(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry, "
"Npart, Massarr, Time, Redshift, FlagSfr, FlagFeedback, Nall, FlagCooling, "
"NumFiles, BoxSize, Omega0, OmegaLambda, HubbleParam, FlagAge, FlagMetals, "
"NallHW, flag_entr_ics)")
self.assertEquals(len(data.gas), 1472)
self.assertEquals(len(data.halo), 0)
self.assertEquals(data.gas[0].key,1)
self.assertEquals(data.gas[1].key,2)
self.assertEquals(data.gas[2].key,3)
self.assertEquals(data.gas[1471].key,1472)
self.assertAlmostRelativeEquals(data.gas[0:5].x,[-0.0713372901082, 0.0713372901082, -0.21178227663, -0.0698266476393, 0.0698266476393] | nbody_system.length, 7)
self.assertAlmostRelativeEquals(data.gas[0:5].u, [0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451, 0.0500000007451] | (nbody_system.length / nbody_system.time)**2, 7 )
self.assertEquals(data.Npart, (1472, 0, 0, 0, 0, 0))
self.assertEquals(data.Time, 0.0)
self.assertEquals(data.Redshift, 0.0)
def test6(self):
print "Testing HDF5 io, with options"
if os.path.exists("test_unit.hdf5"):
os.remove("test_unit.hdf5")
x = datamodel.Particles(2)
x.mass = [1.0, 2.0] | units.kg
io.write_set_to_file(x, "test_unit.hdf5", "hdf5")
x.mass = [10.0, 20.0] | units.kg
io.write_set_to_file(x, "test_unit.hdf5", "hdf5", append_to_file=True)
x.mass = [100.0, 200.0] | units.kg
io.write_set_to_file(x, "test_unit.hdf5", "hdf5")
y = io.read_set_from_file("test_unit.hdf5", "hdf5")
y = y.previous_state()
self.assertAlmostEquals(x.mass, y.mass, 8)
self.assertAlmostEquals([10.0, 20.0] | units.kg,
y.previous_state().mass, 8)
self.assertAlmostEquals([1.0, 2.0] | units.kg,
y.previous_state().previous_state().mass, 8)
self.assertEqual(y.previous_state().previous_state().previous_state(),
None)
io.write_set_to_file(x, "test_unit.hdf5", "hdf5", append_to_file=False)
y = io.read_set_from_file("test_unit.hdf5", "hdf5")
self.assertAlmostEquals(x.mass, y.mass, 8)
self.assertEqual(y.previous_state().previous_state(), None)
os.remove("test_unit.hdf5")
def test13(self):
print "Test convert_gadget_w_to_velocity and return_header for Gadget read_set_from_file"
directory_name = os.path.dirname(__file__)
filename = os.path.join(directory_name, 'tiny_lcdm_data_littleendian.dat')
data = io.read_set_from_file(filename, format='gadget', return_header=False, convert_gadget_w_to_velocity=False) # (default)
self.assertTrue(isinstance(data, tuple))
self.assertEquals(data.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry)")
self.assertEquals(len(data.gas), 32)
self.assertEquals(len(data.halo), 32)
self.assertEquals(data.gas[0].key, 1)
self.assertEquals(data.halo[0].key, 32**3 + 1)
self.assertAlmostRelativeEquals(data.gas[:3].position, [[395.23443604, 395.75210571, 1244.31152344],
[310.17266846, 440.21728516, 2817.06396484], [191.95669556, 465.57223511, 4430.20068359]] | nbody_system.length, 7)
data_converted = io.read_set_from_file(filename, format='gadget', return_header=True, convert_gadget_w_to_velocity=True)
self.assertTrue(isinstance(data_converted, tuple))
self.assertEquals(data_converted.__doc__, "GadgetData(gas, halo, disk, bulge, stars, bndry, "
"Npart, Massarr, Time, Redshift, FlagSfr, FlagFeedback, Nall, FlagCooling, "
"NumFiles, BoxSize, Omega0, OmegaLambda, HubbleParam, FlagAge, FlagMetals, "
"NallHW, flag_entr_ics)")
self.assertEquals(len(data_converted.gas), 32)
self.assertEquals(len(data_converted.halo), 32)
self.assertEquals(data_converted.gas[0].key, 1)
self.assertEquals(data_converted.halo[0].key, 32**3 + 1)
self.assertEquals(data_converted.Npart, (32, 32, 0, 0, 0, 0))
self.assertEquals(data_converted.Time, 1/11.0)
self.assertEquals(data_converted.Redshift, 10.0)
self.assertEquals(data.gas.position, data_converted.gas.position)
self.assertAlmostRelativeEquals(data.gas.velocity, math.sqrt(data_converted.Time) * data_converted.gas.velocity, 7)
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 continue_evolution(sph_code, dynamics_code, t_end, n_steps,
relaxed_giant_output_base_name, do_energy_evolution_plot):
print "Loading snapshots...",
files = os.listdir("snapshots")
files.sort()
files = files[-4:]
print files
binary = read_set_from_file(os.path.join("snapshots", files[0]), format='amuse')
gd_particles = read_set_from_file(os.path.join("snapshots", files[1]), format='amuse')
sph_particles = read_set_from_file(os.path.join("snapshots", files[2]), format='amuse')
snapshotfile = os.path.join("..", "giant_models", relaxed_giant_output_base_name + "_core.amuse")
core_particle = read_set_from_file(snapshotfile, format='amuse')
giant_model = StellarModelInSPH(
gas_particles = sph_particles,
core_particle = gd_particles[0],
core_radius = core_particle.radius)
view_on_giant = Particle()
view_on_giant.position = [0]*3 | units.m
view_on_giant.velocity = [0]*3 | units.m / units.s
print "\nSetting up {0} to simulate inner binary system".format(dynamics_code.__name__)
binary_system = prepare_binary_system(dynamics_code, binary)
print "\nSetting up {0} to simulate giant in SPH".format(sph_code.__name__)
giant_system = prepare_giant_system(sph_code, giant_model, view_on_giant, t_end, n_steps)
print "\nEvolving with bridge between", sph_code.__name__, "and", dynamics_code.__name__
evolve_coupled_system(binary_system, giant_system, t_end, n_steps, do_energy_evolution_plot,
previous_data = os.path.join("snapshots", files[3]))
print "Done"
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.assertAlmostEqual(current_time, y.previous_state().get_timestamp(), 8, in_units=units.Myr)
self.assertAlmostEqual(x.mass, y.mass, 8)
self.assertAlmostEqual(x.position, y.position,8)
self.assertAlmostEqual(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.assertAlmostEqual(current_time, y.previous_state().get_timestamp(), 8, in_units=nbody_system.time)
self.assertAlmostEqual(x.mass, y.mass, 8)
self.assertAlmostEqual(x.position, y.position,8)
self.assertAlmostEqual(x.velocity, y.velocity,8)
os.remove("time_test_unit.tsf")
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 continue_evolution(sph_code, dynamics_code, t_end, n_steps,
relaxed_giant_output_base_name, do_energy_evolution_plot):
print("Loading snapshots...", end=' ')
files = os.listdir("snapshots")
files.sort()
files = files[-4:]
print(files)
binary = read_set_from_file(os.path.join("snapshots", files[0]), format='amuse')
gd_particles = read_set_from_file(os.path.join("snapshots", files[1]), format='amuse')
sph_particles = read_set_from_file(os.path.join("snapshots", files[2]), format='amuse')
snapshotfile = os.path.join("..", "giant_models", relaxed_giant_output_base_name + "_core.amuse")
core_particle = read_set_from_file(snapshotfile, format='amuse')
giant_model = StellarModelInSPH(
gas_particles = sph_particles,
core_particle = gd_particles[0],
core_radius = core_particle.radius)
view_on_giant = Particle()
view_on_giant.position = [0]*3 | units.m
view_on_giant.velocity = [0]*3 | units.m / units.s
print("\nSetting up {0} to simulate inner binary system".format(dynamics_code.__name__))
binary_system = prepare_binary_system(dynamics_code, binary)
print("\nSetting up {0} to simulate giant in SPH".format(sph_code.__name__))
giant_system = prepare_giant_system(sph_code, giant_model, view_on_giant, t_end, n_steps)
print("\nEvolving with bridge between", sph_code.__name__, "and", dynamics_code.__name__)
evolve_coupled_system(binary_system, giant_system, t_end, n_steps, do_energy_evolution_plot,
previous_data = os.path.join("snapshots", files[3]))
print("Done")
def load_young_star_cluster(filename_base):
filename_base = os.path.join("..", "..", "initial_conditions", filename_base)
with open(filename_base + "info.pkl", "rb") as infile:
converter = cPickle.load(infile)[0]
stars = read_set_from_file(filename_base + "stars.amuse", "amuse")
gas = read_set_from_file(filename_base + "gas.amuse", "amuse")
# gas = read_set_from_file(filename_base + "gas_relaxed.amuse", "amuse")
return stars, gas, converter
def test7(self):
directory = os.path.dirname(__file__)
set, converter = io.read_set_from_file(os.path.join(directory, 'evolved.dyn'), 'starlab', return_converter = True)
self.assertEquals(len(set), 20)
unconverted_set = io.read_set_from_file(os.path.join(directory, 'evolved.dyn'), 'starlab', must_scale = False)
self.assertEquals(len(unconverted_set), 20)
self.assertAlmostRelativeEquals(converter.to_nbody(set.x), unconverted_set.x)
def load_giant_model(file_base_name):
snapshotfile = os.path.join("..", "giant_models", file_base_name + "_gas.amuse")
sph_particles = read_set_from_file(snapshotfile, format='amuse')
snapshotfile = os.path.join("..", "giant_models", file_base_name + "_core.amuse")
gd_particles = read_set_from_file(snapshotfile, format='amuse')
core_radius = gd_particles[0].radius
gd_particles[0].radius = 0.0 | units.m
giant_model = StellarModelInSPH(
gas_particles = sph_particles,
core_particle = gd_particles[0],
core_radius = core_radius)
return giant_model
def load_giant_model(file_base_name):
snapshotfile = os.path.join("..", "giant_models", file_base_name + "_gas.amuse")
sph_particles = read_set_from_file(snapshotfile, format='amuse')
snapshotfile = os.path.join("..", "giant_models", file_base_name + "_core.amuse")
gd_particles = read_set_from_file(snapshotfile, format='amuse')
core_radius = gd_particles[0].radius
gd_particles[0].radius = 0.0 | units.m
giant_model = StellarModelInSPH(
gas_particles = sph_particles,
core_particle = gd_particles[0],
core_radius = core_radius)
return giant_model
def plot_densities(self):
print "Plotting x-density as function of time"
if not os.path.exists('out/{0}/plots/density'.format(self.timestamp)):
os.mkdir('out/{0}/plots/density'.format(self.timestamp))
density_gasA_list = [] | (units.MSun / units.Mpc**3)
density_dmA_list = [] | (units.MSun / units.Mpc**3)
density_gasB_list = [] | (units.MSun / units.Mpc**3)
density_dmB_list = [] | (units.MSun / units.Mpc**3)
time_list = [] | units.Gyr
paths = sorted(glob.glob('out/{0}/data/gasA_*'.format(self.timestamp)),
key=os.path.getmtime)
tot = len(paths) - 1
for i, path in enumerate(paths):
time = path[28:-6] # fix due to round error when storing the data.
print_progressbar(i, tot)
gasA = read_set_from_file(
"out/{0}/data/gasA_{1}.amuse".format(self.timestamp, time),
"amuse")
dmA = read_set_from_file(
"out/{0}/data/dmA_{1}.amuse".format(self.timestamp, time),
"amuse")
gasB = read_set_from_file(
"out/{0}/data/gasB_{1}.amuse".format(self.timestamp, time),
"amuse")
dmB = read_set_from_file(
"out/{0}/data/dmB_{1}.amuse".format(self.timestamp, time),
"amuse")
density_gasA = gasA.mass / (4. / 3 * numpy.pi * gasA.h_smooth**3)
density_dmA = dmA.mass / (4. / 3 * numpy.pi * dmA.radius**3)
density_gasB = gasB.mass / (4. / 3 * numpy.pi * gasB.h_smooth**3)
density_dmB = dmA.mass / (4. / 3 * numpy.pi * dmA.radius**3)
fig = pyplot.figure(figsize=(12, 10), dpi=50)
plot(gasA.x, density_gasA, label="gasA", c="r", ls="solid")
plot(dmA.x, density_dmA, label="dmA", c="r", ls="dashed")
plot(gasB.x, density_gasB, label="gasB", c="g", ls="solid")
plot(dmB.x, density_dmB, label="dmB", c="g", ls="dashed")
xlabel("x")
ylabel("Density")
pyplot.legend(bbox_to_anchor=(1.05, 1), loc=2, borderaxespad=0.)
pyplot.show()
pyplot.savefig('out/{0}/plots/density/density_{1}'.format(
self.timestamp, time),
dpi=50)
pyplot.close()
if i > 10:
break
print
def read_state_from_file(restart_file, gravity_code,
resolve_collision_code_creation_function, kep):
stars = io.read_set_from_file(restart_file + ".stars.hdf5",
'hdf5',
version='2.0').copy()
stars_python = io.read_set_from_file(restart_file + ".stars_python.hdf5",
'hdf5',
version='2.0').copy()
with open(restart_file + ".bookkeeping", "rb") as f:
bookkeeping = pickle.load(f)
if os.path.isfile(restart_file + ".extra_stuff"):
with open(restart_file + ".extra_stuff", "rb") as f:
extra_stuff = pickle.load(f)
else:
extra_stuff = None
print bookkeeping
root_to_tree = {}
for root in stars:
if hasattr(root, 'components') and not root.components is None:
root_to_tree[root] = datamodel.trees.BinaryTreeOnParticle(
root.components[0])
gravity_code.set_begin_time(bookkeeping['model_time'])
gravity_code.particles.add_particles(stars)
gravity_code.commit_particles()
multiples_code = multiples.Multiples(
gravity_code, resolve_collision_code_creation_function, kep)
#multiples_code.neighbor_distance_factor = 1.0
#multiples_code.neighbor_veto = False
#multiples_code.neighbor_distance_factor = 2.0
#multiples_code.neighbor_veto = True
multiples_code.neighbor_distance_factor = bookkeeping[
'neighbor_distance_factor']
multiples_code.neighbor_veto = bookkeeping['neighbor_veto']
multiples_code.multiples_external_tidal_correction = bookkeeping[
'multiples_external_tidal_correction']
multiples_code.multiples_integration_energy_error = bookkeeping[
'multiples_integration_energy_error']
multiples_code.multiples_internal_tidal_correction = bookkeeping[
'multiples_internal_tidal_correction']
multiples.root_index = bookkeeping['root_index']
multiples_code.root_to_tree = root_to_tree
#multiples_code.set_model_time = bookkeeping['model_time']
with open(restart_file + ".conf", "rb") as f:
config = pickle.load(f)
random.setstate(pickle.loads(config["py_seed"]))
numpy.random.set_state(pickle.loads(config["numpy_seed"]))
return stars_python, bookkeeping['model_time'], multiples_code, extra_stuff
def read_state_from_file(restart_file, gravity_code, kep, MT=0):
# Function to load from file. If you change params in
# write_state_to_file, make sure you match the changes here.
stars = io.read_set_from_file(restart_file + ".stars.hdf5",
'amuse',
version='2.0',
close_file=True).copy()
stars_python = io.read_set_from_file(restart_file + ".stars_python.hdf5",
'amuse',
version='2.0',
close_file=True).copy()
with open(restart_file + ".params", "rb") as f:
params = pickle.load(f)
f.close()
if MT == 0:
root_to_tree = {}
for root in stars:
if hasattr(root, 'components') and not root.components is None:
root_to_tree[root] \
= datamodel.trees.BinaryTreeOnParticle(root.components[0])
else:
root_to_tree = MT
gravity_code.parameters.timestep_parameter = params['timestep_parameter']
gravity_code.parameters.epsilon_squared = params['epsilon_squared']
gravity_code.particles.add_particles(stars)
gravity_code.commit_particles() # sets time steps
multiples_code = multiples.Multiples(gravity_code, new_smalln, kep)
multiples_code.neighbor_veto = params['neighbor_veto']
multiples_code.multiples_external_tidal_correction \
= params['multiples_external_tidal_correction']
multiples_code.multiples_integration_energy_error \
= params['multiples_integration_energy_error']
multiples_code.multiples_internal_tidal_correction \
= params['multiples_internal_tidal_correction']
multiples.root_index = params['root_index']
multiples_code.root_to_tree = root_to_tree
print("\nread state from file ", restart_file, \
'at time', params['model_time'])
return stars_python, params['model_time'], params['delta_t'], \
params['EZero'], params['CPUZero'], multiples_code
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 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 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 test3(self):
filename = os.path.join(os.path.dirname(__file__), 'FinalSnapshot.out')
set = io.read_set_from_file(filename, 'dyn')
print set
print "set[0].parent.mass", set[0].parent.mass
self.assertAlmostRelativeEquals(0.000227826766314251919 * 617.75586357299929284496, set[9].mass.value_in(units.MSun), 12)
self.assertAlmostRelativeEquals(617.75586357299929284496 * 0.953575109205781479, set[0].parent.mass.value_in(units.MSun), 12)
def plots(i):
g = read_set_from_file('dump-%6.6i' % i, 'amuse')
r = ((g.x**2+g.y**2+g.z**2)**0.5).value_in(units.kpc)
v = ((g.vx**2+g.vy**2+g.vz**2)**0.5).value_in(units.kms)
cs = (g.u**0.5).value_in(units.kms)
xion = xion_from_u(g.u).value_in(units.none)
rho = 3./4/numpy.pi*8*g.mass/g.radius**3
dens = (rho).value_in(units.amu/units.cm**3)
pres = (g.u*rho).value_in(units.g/units.cm/units.s**2)
mach = v/cs
pyplot.figure(figsize=(6, 6))
pyplot.semilogy(r/15, xion, 'r .')
pyplot.semilogy(r/15, 1-xion, 'g .')
pyplot.xlim((0., 1.))
pyplot.ylim((1.e-6, 1.))
pyplot.savefig('xion-part-%6.6i.png' % i)
aplot(i, 'xion', ((r/15, xion, 'r'), (r/15, 1-xion, 'g')),
xlim=(0., 1.), ylim=(1.e-6, 1.))
aplot(i, 'pres', ((r/15, pres, 'r'),),
xlim=(0., 1.), ylim=(1.e-17, 1.e-14))
aplot(i, 'rho', ((r/15, dens, 'r'),),
xlim=(0., 1.), ylim=(0.0001, 0.01))
aplot(i, 'mach', ((r/15, mach, 'r'),),
xlim=(0., 1.), ylim=(1.e-5, 10.))
def test15(self):
test_results_path = self.get_path_to_results()
output_files = []
filebase = os.path.join(test_results_path, "test_distributed_sets")
for i in [0, 1]:
output_file = filebase + "_%6.6i" % i
output_files.append(output_file)
if os.path.exists(output_file):
os.remove(output_file)
x = DistributedParticles(size=8, number_of_workers=2)
x.index = [0, 10, 10, 10, 0, 0, 10, 0]
x.mass = [1, 2, 3, 4, 5, 6, 7, 8] | units.MSun
highx = x.select_array(select_example_function, ("index", ))
files = highx.write_set_to_file(filebase, "amuse")
self.assertEqual(files, output_files)
expected_local_sizes = [3, 1]
for i, f in enumerate(files):
self.assertTrue(os.path.isfile(f))
p = read_set_from_file(f, "amuse")
self.assertEqual(len(p), expected_local_sizes[i])
self.assertEqual(p.collection_attributes.global_size, 4)
self.assertEqual(p.collection_attributes.local_size,
expected_local_sizes[i])
def read_set_from_file(self, reference_in, filenames_in, fileformat):
ceil = lambda x, y: (x / y + (x % y > 0))
filenames = decode_and_load(filenames_in)
if reference_in != 0:
raise Exception("reference_in hould be zero here!")
real_particles, local_offset, local_size, global_size = self.references_to_particles[
reference_in]
start = self.rank * ceil(len(filenames), self.number_of_processes)
end = (self.rank + 1) * ceil(len(filenames), self.number_of_processes)
for i in range(start, end):
if i < len(filenames):
p = read_set_from_file(filenames[i], fileformat)
real_particles.add_particles(p)
input = numpy.zeros(1, dtype='int64')
output = numpy.zeros(self.number_of_processes, dtype='int64')
input[0] = len(real_particles)
self.mpi_comm.Allgather([input, MPI.INTEGER], [output, MPI.INTEGER])
total_size = 0
local_offset = 0
for i, current_size in enumerate(output):
if i < self.rank:
local_offset += current_size
total_size += current_size
referenced_particles = ReferencedParticles(real_particles,
local_offset,
len(real_particles),
total_size)
self.references_to_particles[reference_in] = referenced_particles
return 0
def read_set_from_file(self, reference_in, filenames_in, fileformat):
ceil = lambda x, y: (x / y + (x % y > 0))
filenames = decode_and_load(filenames_in)
if reference_in != 0:
raise Exception("reference_in hould be zero here!")
real_particles, local_offset, local_size, global_size = self.references_to_particles[reference_in]
start = self.rank * ceil(len(filenames), self.number_of_processes)
end = (self.rank + 1) * ceil(len(filenames), self.number_of_processes)
for i in range(start, end):
if i < len(filenames):
p = read_set_from_file(filenames[i], fileformat)
real_particles.add_particles(p)
input = numpy.zeros(1, dtype="int64")
output = numpy.zeros(self.number_of_processes, dtype="int64")
input[0] = len(real_particles)
self.mpi_comm.Allgather([input, MPI.INTEGER], [output, MPI.INTEGER])
total_size = 0
local_offset = 0
for i, current_size in enumerate(output):
if i < self.rank:
local_offset += current_size
total_size += current_size
referenced_particles = ReferencedParticles(real_particles, local_offset, len(real_particles), total_size)
self.references_to_particles[reference_in] = referenced_particles
return 0
def test54(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)
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[0][0]
stars[1].gas = gas[1][0]
self.assertAlmostRelativeEquals(stars[0].gas.y, 1.0 | units.km)
self.assertAlmostRelativeEquals(stars[1].gas.y, 4.0 | units.km)
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.y, 1.0 | units.km)
self.assertAlmostRelativeEquals(loaded_stars[1].gas.y, 4.0 | units.km)
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.assertEqual(output.new_attribute, 2 * output.mass)
output = next(output.iter_history())
output.new_attribute = 2 * output.mass
self.assertEqual(output.new_attribute, 2 * output.mass)
def test52(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test52" + 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 = Particles(1)
gas[0].y = 1.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.assertEqual(loaded_stars[0].gas[0].key, gas[0].key)
self.assertAlmostRelativeEquals(loaded_stars[0].gas[0].y,
1.0 | units.km)
self.assertEqual(loaded_stars[1].gas, None)
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 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 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 test13(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test13" + 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, 2, 3], [4, 5, 6]] | units.km
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.assertEqual(loaded[0].md, [1, 2, 3] | units.km)
self.assertEqual(loaded[1].md, [4, 5, 6] | units.km)
self.assertEqual(loaded.md[0], [1, 2, 3] | units.km)
self.assertEqual(loaded.md[1], [4, 5, 6] | units.km)
self.assertEqual(
self.get_version_in_store(loaded)[0].md, [1, 2, 3] | units.km)
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 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 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 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 test4(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test4" + self.store_version() + ".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
particles = Particles(10)
particles.mass = 1.0 | units.kg
particles.x = 2.0 | units.m
x = particles.savepoint(2.0 | units.s)
io.write_set_to_file(x,
output_file,
format='amuse',
version=self.store_version())
particles_from_file = io.read_set_from_file(
output_file,
format='amuse',
version=self.store_version(),
close_file=True)
particles_in_memory = particles_from_file.copy()
self.assertAlmostRelativeEquals(particles_in_memory.mass[1],
1.0 | units.kg)
particles_in_memory.savepoint(4.0 | units.s)
self.assertAlmostRelativeEquals(particles_from_file.mass[2],
1.0 | units.kg)
def test55(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test25" + self.store_version() + ".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
particles1 = Particles(2)
particles1[0].x = 1.0 | units.km
particles1[1].x = 2.0 | units.km
particles2 = Particles(2)
particles2[0].x = 3.0 | units.km
particles2[1].x = 4.0 | units.km
particles_superset = ParticlesSuperset([particles1, particles2])
self.assertAlmostRelativeEquals(particles_superset[0].x,
1.0 | units.km)
self.assertAlmostRelativeEquals(particles_superset[2].x,
3.0 | units.km)
io.write_set_to_file(particles_superset,
output_file,
"hdf5",
version=self.store_version())
loaded_particles_superset = io.read_set_from_file(
output_file, "hdf5", version=self.store_version())
self.assertAlmostRelativeEquals(loaded_particles_superset[0].x,
1.0 | units.km)
self.assertAlmostRelativeEquals(loaded_particles_superset[2].x,
3.0 | units.km)
def test4(self):
directory_name = os.path.dirname(__file__)
filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
x = gadget.GadgetFileFormatProcessor()
result = io.read_set_from_file(filename, format='gadget')
self.assertEquals(len(result[0]), 1472)
self.assertEquals(len(result[1]), 0)
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 test15(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test15" + 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())
processor = self.store_factory()(output_file,
True,
open_for_writing=True)
loaded = processor.load()
self.get_version_in_store(loaded)[0].md = [[3, 1], [3, 4], [5, 2]]
self.assertEqual(
self.get_version_in_store(loaded)[0].md, [[3, 1], [3, 4], [5, 2]])
processor.close()
loaded = io.read_set_from_file(output_file,
"hdf5",
version=self.store_version())
self.assertEqual(loaded[0].md, [[3, 1], [3, 4], [5, 2]])
self.assertEqual(loaded[1].md, [[4, 6], [5, 7], [6, 8]])
def test3(self):
directory_name = os.path.dirname(__file__)
filename = os.path.join(directory_name, 'gadget_snapshot')
x = gadget.GadgetFileFormatProcessor()
result = io.read_set_from_file(filename, format='gadget')
self.assertEquals(len(result[0]), 1000)
self.assertEquals(len(result[1]), 10000)
def test21(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test21" + 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",
close_file=True,
version=self.store_version(),
copy_history=False)
self.assertEqual(loaded[0].md, [[1, 3], [2, 4], [3, 5]])
self.assertEqual(loaded[1].md, [[4, 6], [5, 7], [6, 8]])
previous = loaded.previous_state()
self.assertEqual(previous, None)
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.assertEqual(loaded[0].x, 1.0 | units.km)
self.assertEqual(loaded[1].y, 2.0 | units.km)
def test24(self):
test_results_path = self.get_path_to_results()
output_file = os.path.join(test_results_path,
"test24" + self.store_version() + ".hdf5")
if os.path.exists(output_file):
os.remove(output_file)
particles = Particles(20)
#~ particles = Particles(1000000) # For testing memory usage
particles.mass = 1.0 | units.kg
particles.position = [0.0, 0.0, 0.0] | units.m
for i in range(10):
particles.position += [1.0, 2.0, 3.0] | units.m
io.write_set_to_file(particles,
output_file,
format='amuse',
version=self.store_version(),
append_to_file=True)
particles_from_file = io.read_set_from_file(
output_file, format='amuse', version=self.store_version())
os.remove(output_file)
self.assertEqual(len(list(particles_from_file.history)), 10)
for i, snap in enumerate(particles_from_file.history):
self.assertEqual(len(snap), 20)
self.assertEqual((i + 1) * ([1.0, 2.0, 3.0] | units.m),
snap.center_of_mass())
def read_state_from_file(restart_file, gravity_code, kep, MT=0):
# Function to load from file. If you change params in
# write_state_to_file, make sure you match the changes here.
stars = io.read_set_from_file(restart_file+".stars.hdf5",
'amuse', version='2.0',
close_file=True).copy()
stars_python = io.read_set_from_file(
restart_file+".stars_python.hdf5",
'amuse',version='2.0',
close_file=True).copy()
with open(restart_file + ".params", "rb") as f:
params = pickle.load(f)
f.close()
if MT == 0:
root_to_tree = {}
for root in stars:
if hasattr(root, 'components') and not root.components is None:
root_to_tree[root] \
= datamodel.trees.BinaryTreeOnParticle(root.components[0])
else:
root_to_tree = MT
gravity_code.parameters.timestep_parameter = params['timestep_parameter']
gravity_code.parameters.epsilon_squared = params['epsilon_squared']
gravity_code.particles.add_particles(stars)
gravity_code.commit_particles() # sets time steps
multiples_code = multiples.Multiples(gravity_code, new_smalln, kep)
multiples_code.neighbor_veto = params['neighbor_veto']
multiples_code.multiples_external_tidal_correction \
= params['multiples_external_tidal_correction']
multiples_code.multiples_integration_energy_error \
= params['multiples_integration_energy_error']
multiples_code.multiples_internal_tidal_correction \
= params['multiples_internal_tidal_correction']
multiples.root_index = params['root_index']
multiples_code.root_to_tree = root_to_tree
print "\nread state from file ", restart_file, \
'at time', params['model_time']
return stars_python, params['model_time'], params['delta_t'], \
params['EZero'], params['CPUZero'], multiples_code
def test9(self):
filename = os.path.join(os.path.dirname(__file__), 'plummer128.nemo')
particles = io.read_set_from_file(filename, format="nemobin")
self.assertEquals(len(particles), 128)
self.assertAlmostEquals(particles.total_mass(), 1.0 | nbody_system.mass)
self.assertAlmostEquals(particles.center_of_mass(), 0.0 | nbody_system.length)
self.assertAlmostEquals(particles.center_of_mass_velocity(), 0.0 | nbody_system.speed)
self.assertAlmostEquals(particles.kinetic_energy(), 0.230214395174 | nbody_system.energy)
def main(filename):
pyplot.figure(figsize=(12,12))
particles = read_set_from_file(filename, "hdf5")
for si in particles.history:
scatter(si.x, si.y, s=100)
xlabel("x")
ylabel("y")
pyplot.show()
def read_gas_file(self,filename):
p=read_set_from_file(filename,'amuse')
mass=p.mass.number
hsml=p.smoothing_length.number
x=p.x.number
y=p.y.number
z=p.z.number
rho=p.rho.number
u=p.internal_energy.number
xe=numpy.zeros_like(x)
return mass, hsml, x, y, z, rho, xe, u
def test6(self):
directory_name = os.path.dirname(__file__)
filename = os.path.join(directory_name, 'gassphere_littleendian.dat')
x = gadget.GadgetFileFormatProcessor()
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)
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 read_gas_file(self,filename):
p=read_set_from_file(filename,'amuse')
mass=p.mass.number
hsml=p.smoothing_length.number
x=p.x.number
y=p.y.number
z=p.z.number
rho=p.rho.number
u=p.internal_energy.number
xe=numpy.zeros_like(x)
return create_particle_set(mass=mass | (10**10*units.MSun), h_smooth=hsml | (units.kpc),
x=x | (units.kpc), y=y| (units.kpc), z=z| (units.kpc), rho=rho | ((10**10*units.MSun) /(units.kpc)**3),
xion=xe, u=u| (10**5 * units.cm/units.s)**2)
def parse_file(f1): """ Parses two files, each using a different coordinate system (one for each star) Note: Data stored in a different format will require a different function f1 = coordinates with respect to center of mass snapshot_dir = directory of choice """ #print " ** plotting: " bodies = read_set_from_file(snapshot_dir + "/" + f1, 'hdf5') return bodies
def make_plots(Ncl, Rcl, t_end):
""" Not finished """
try:
import matplotlib
matplotlib.use("Agg")
from matplotlib import pyplot
except ImportError:
print "Unable to produce plots: couldn't find matplotlib"
else:
fig = pyplot.figure(figsize=[10, 10])
for data in ["hybrid", "tree", "direct"]
print "Generating plot data of {0} run".format(data)
data_file = "CA_Exam_TLRH_{0}.amuse".format(data)
stars_below_cut, stars_above_cut =\
read_set_from_file(data_file, 'amuse',
names=("stars_below_cut", "stars_below_cut"))
lim = abs(next(stars_below_cut.history).position).max()
if data == "hybrid":
pass
# dashed curve
# points = aplot.scatter(dE, time, c='orange',
# label='Hybrid')
if data == "direct":
pass
# dotted curve
# points = aplot.scatter(dE, time, c='red',
# label='Direct')
if data == "tree":
pass
# solid curve
# points = aplot.scatter(dE, time, c='green',
# label='Tree')
aplot.xlabel("Time")
aplot.ylabel("Relative Energy Error")
pyplot.axis('equal')
aplot.xlim(-lim, lim)
aplot.ylim(-lim, lim)
pyplot.legend()
pyplot.title(r'Cluster with $N=${0}, $r=${1}'
.format(Ncl, Rcl ) +
r', evolved until $t_{\rm end}$={0}'
.format(t_end))
pyplot.savefig("out.png")
pyplot.close()
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")
