def test2(self):
print "Test basic particle attributes and scale_to_standard - SI units"
convert_nbody = nbody_system.nbody_to_si(1 | units.MSun, 1 | units.parsec)
particles = Particles(2)
particles.position = [[-1, 0, 0], [1,0,0]] | units.parsec
particles.velocity = [[-1, 0, 0], [1,0,0]] | units.parsec / units.Myr
particles.mass = 0.5 | units.MSun
self.assertAlmostRelativeEquals(particles.total_mass(), 1.0 | units.MSun)
self.assertAlmostRelativeEquals(particles.kinetic_energy(), 1.0 * (0.5 | units.MSun) * (1 |units.parsec / units.Myr) **2 )
self.assertAlmostRelativeEquals(particles.potential_energy(), -constants.G * (0.5 | units.MSun) ** 2 / ([2,0,0] | units.parsec).length() )
self.assertAlmostRelativeEquals(particles.virial_radius(), 4.0 | units.parsec)
particles.scale_to_standard(convert_nbody)
self.assertAlmostRelativeEquals(particles.total_mass(), convert_nbody.to_si(1.0 | nbody_system.mass))
self.assertAlmostRelativeEquals(particles.kinetic_energy(), convert_nbody.to_si(0.25 | nbody_system.energy))
self.assertAlmostRelativeEquals(particles.potential_energy().as_quantity_in(units.J), convert_nbody.to_si(-0.5 | nbody_system.energy).as_quantity_in(units.J), 12)
self.assertAlmostRelativeEquals(particles.virial_radius(), convert_nbody.to_si(1.0 | nbody_system.length))
particles.scale_to_standard(convert_nbody, virial_ratio=1) # unbound
self.assertAlmostRelativeEquals(particles.kinetic_energy(), 0.5 * constants.G * (1 | units.MSun**2 / units.parsec), 13)
self.assertAlmostRelativeEquals(particles.potential_energy(), -0.5 * constants.G * (1 | units.MSun**2 / units.parsec))
particles.scale_to_standard(convert_nbody, virial_ratio=0) # velocities zeroed
self.assertAlmostRelativeEquals(particles.kinetic_energy(), 0 | units.J)
self.assertAlmostRelativeEquals(particles.potential_energy(), -0.5 * constants.G * (1 | units.MSun**2 / units.parsec))
def test2(self):
print(
"Test basic particle attributes and scale_to_standard - SI units")
convert_nbody = nbody_system.nbody_to_si(1 | units.MSun,
1 | units.parsec)
particles = Particles(2)
particles.position = [[-1, 0, 0], [1, 0, 0]] | units.parsec
particles.velocity = [[-1, 0, 0], [1, 0, 0]] | units.parsec / units.Myr
particles.mass = 0.5 | units.MSun
self.assertAlmostRelativeEquals(particles.total_mass(),
1.0 | units.MSun)
self.assertAlmostRelativeEquals(
particles.kinetic_energy(),
1.0 * (0.5 | units.MSun) * (1 | units.parsec / units.Myr)**2)
self.assertAlmostRelativeEquals(
particles.potential_energy(), -constants.G *
(0.5 | units.MSun)**2 / ([2, 0, 0] | units.parsec).length())
self.assertAlmostRelativeEquals(particles.virial_radius(),
4.0 | units.parsec)
particles.scale_to_standard(convert_nbody)
self.assertAlmostRelativeEquals(
particles.total_mass(),
convert_nbody.to_si(1.0 | nbody_system.mass))
self.assertAlmostRelativeEquals(
particles.kinetic_energy(),
convert_nbody.to_si(0.25 | nbody_system.energy))
self.assertAlmostRelativeEquals(
particles.potential_energy().as_quantity_in(units.J),
convert_nbody.to_si(-0.5 | nbody_system.energy).as_quantity_in(
units.J), 12)
self.assertAlmostRelativeEquals(
particles.virial_radius(),
convert_nbody.to_si(1.0 | nbody_system.length))
particles.scale_to_standard(convert_nbody, virial_ratio=1) # unbound
self.assertAlmostRelativeEquals(
particles.kinetic_energy(),
0.5 * constants.G * (1 | units.MSun**2 / units.parsec), 13)
self.assertAlmostRelativeEquals(
particles.potential_energy(),
-0.5 * constants.G * (1 | units.MSun**2 / units.parsec))
particles.scale_to_standard(convert_nbody,
virial_ratio=0) # velocities zeroed
self.assertAlmostRelativeEquals(particles.kinetic_energy(),
0 | units.J)
self.assertAlmostRelativeEquals(
particles.potential_energy(),
-0.5 * constants.G * (1 | units.MSun**2 / units.parsec))
def test1(self):
print(
"Test basic particle attributes and scale_to_standard - nbody units"
)
particles = Particles(2)
particles.position = [[-1, 0, 0], [1, 0, 0]] | nbody_system.length
particles.velocity = [[-1, 0, 0], [1, 0, 0]
] | nbody_system.length / nbody_system.time
particles.mass = 0.4 | nbody_system.mass
self.assertAlmostRelativeEquals(particles.total_mass(),
0.8 | nbody_system.mass)
self.assertAlmostRelativeEquals(particles.kinetic_energy(),
0.4 | nbody_system.energy)
self.assertAlmostRelativeEquals(
particles.potential_energy(G=nbody_system.G),
-0.08 | nbody_system.energy)
self.assertAlmostRelativeEquals(particles.virial_radius(),
4.0 | nbody_system.length)
particles.scale_to_standard()
self.assertAlmostRelativeEquals(particles.total_mass(),
1.0 | nbody_system.mass)
self.assertAlmostRelativeEquals(particles.kinetic_energy(),
0.25 | nbody_system.energy)
self.assertAlmostRelativeEquals(
particles.potential_energy(G=nbody_system.G),
-0.5 | nbody_system.energy)
self.assertAlmostRelativeEquals(particles.virial_radius(),
1.0 | nbody_system.length)
particles.scale_to_standard(virial_ratio=1) # unbound
self.assertAlmostRelativeEquals(particles.kinetic_energy(),
0.5 | nbody_system.energy)
self.assertAlmostRelativeEquals(
particles.potential_energy(G=nbody_system.G),
-0.5 | nbody_system.energy)
particles.scale_to_standard(virial_ratio=0) # velocities zeroed
self.assertAlmostRelativeEquals(particles.kinetic_energy(),
0 | nbody_system.energy)
self.assertAlmostRelativeEquals(
particles.potential_energy(G=nbody_system.G),
-0.5 | nbody_system.energy)
def test1(self):
print "Test basic particle attributes and scale_to_standard - nbody units"
particles = Particles(2)
particles.position = [[-1, 0, 0], [1,0,0]] | nbody_system.length
particles.velocity = [[-1, 0, 0], [1,0,0]] | nbody_system.length/nbody_system.time
particles.mass = 0.4 | nbody_system.mass
self.assertAlmostRelativeEquals(particles.total_mass(), 0.8 | nbody_system.mass)
self.assertAlmostRelativeEquals(particles.kinetic_energy(), 0.4 | nbody_system.energy)
self.assertAlmostRelativeEquals(particles.potential_energy(G=nbody_system.G), -0.08 | nbody_system.energy)
self.assertAlmostRelativeEquals(particles.virial_radius(), 4.0 | nbody_system.length)
particles.scale_to_standard()
self.assertAlmostRelativeEquals(particles.total_mass(), 1.0 | nbody_system.mass)
self.assertAlmostRelativeEquals(particles.kinetic_energy(), 0.25 | nbody_system.energy)
self.assertAlmostRelativeEquals(particles.potential_energy(G=nbody_system.G), -0.5 | nbody_system.energy)
self.assertAlmostRelativeEquals(particles.virial_radius(), 1.0 | nbody_system.length)
particles.scale_to_standard(virial_ratio=1) # unbound
self.assertAlmostRelativeEquals(particles.kinetic_energy(), 0.5 | nbody_system.energy)
self.assertAlmostRelativeEquals(particles.potential_energy(G=nbody_system.G), -0.5 | nbody_system.energy)
particles.scale_to_standard(virial_ratio=0) # velocities zeroed
self.assertAlmostRelativeEquals(particles.kinetic_energy(), 0 | nbody_system.energy)
self.assertAlmostRelativeEquals(particles.potential_energy(G=nbody_system.G), -0.5 | nbody_system.energy)
def test2(self):
colliders = Particles(2)
colliders.mass = [5, 5] | units.kg
colliders.position = [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]] | units.m
colliders.velocity = [[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]] | units.m / units.s
merged = StickySpheres(mass_loss=0.2).handle_collision(colliders[0], colliders[1])
self.assertTrue(isinstance(merged, Particles))
self.assertEqual(merged.mass, 8 | units.kg)
self.assertAlmostEqual(merged.position, [0.5, 0.5, 0.5] | units.m)
self.assertAlmostEqual(merged.velocity, [1.0, 1.0, 1.0] | units.m / units.s)
copy = colliders.copy()
copy.move_to_center()
self.assertAlmostEqual(colliders.kinetic_energy(), merged.as_set().kinetic_energy() / 0.8 + copy.kinetic_energy())
def test2(self):
colliders = Particles(2)
colliders.mass = [5, 5] | units.kg
colliders.position = [[0.0, 0.0, 0.0], [1.0, 1.0, 1.0]] | units.m
colliders.velocity = [[0.0, 0.0, 0.0], [2.0, 2.0, 2.0]] | units.m / units.s
merged = StickySpheres(mass_loss=0.2).handle_collision(colliders[0], colliders[1])
self.assertTrue(isinstance(merged, Particles))
self.assertEqual(merged.mass, 8 | units.kg)
self.assertAlmostEqual(merged.position, [0.5, 0.5, 0.5] | units.m)
self.assertAlmostEqual(merged.velocity, [1.0, 1.0, 1.0] | units.m / units.s)
copy = colliders.copy()
copy.move_to_center()
self.assertAlmostEqual(colliders.kinetic_energy(), merged.as_set().kinetic_energy() / 0.8 + copy.kinetic_energy())
def test1(self):
colliders = Particles(2)
colliders.mass = [5, 2] | units.kg
colliders.position = [[0.0, 0.0, 0.0], [0.7, 1.4, -0.35]] | units.m
colliders.velocity = [[0.4, -0.6, 0.0], [0.0, 0.0, -3.0]] | units.m / units.s
self.assertAlmostEqual(colliders.center_of_mass_velocity().length(), 1.0 | units.m / units.s)
merged = StickySpheres().handle_collision(colliders[0], colliders[1])
self.assertTrue(isinstance(merged, Particles))
self.assertEqual(merged.mass, 7 | units.kg)
self.assertAlmostEqual(merged.position, [0.2, 0.4, -0.1] | units.m)
self.assertAlmostEqual(merged.velocity, ([2.0, -3.0, -6.0] | units.m / units.s) / 7.0)
self.assertAlmostEqual(merged.velocity.length(), 1.0 | units.m / units.s)
copy = colliders.copy()
copy.move_to_center()
self.assertAlmostEqual(colliders.kinetic_energy(), merged.as_set().kinetic_energy() + copy.kinetic_energy())
def test1(self):
colliders = Particles(2)
colliders.mass = [5, 2] | units.kg
colliders.position = [[0.0, 0.0, 0.0], [0.7, 1.4, -0.35]] | units.m
colliders.velocity = [[0.4, -0.6, 0.0], [0.0, 0.0, -3.0]] | units.m / units.s
self.assertAlmostEqual(colliders.center_of_mass_velocity().length(), 1.0 | units.m / units.s)
merged = StickySpheres().handle_collision(colliders[0], colliders[1])
self.assertTrue(isinstance(merged, Particles))
self.assertEqual(merged.mass, 7 | units.kg)
self.assertAlmostEqual(merged.position, [0.2, 0.4, -0.1] | units.m)
self.assertAlmostEqual(merged.velocity, ([2.0, -3.0, -6.0] | units.m / units.s) / 7.0)
self.assertAlmostEqual(merged.velocity.length(), 1.0 | units.m / units.s)
copy = colliders.copy()
copy.move_to_center()
self.assertAlmostEqual(colliders.kinetic_energy(), merged.as_set().kinetic_energy() + copy.kinetic_energy())
class GravityCodeForTesting(object):
def __init__(self):
self.particles = Particles()
self.model_time = quantities.zero
def evolve_model(self, t_end):
self.particles.position += self.particles.velocity * (t_end -
self.model_time)
self.model_time = t_end
@property
def potential_energy(self):
G = nbody_system.G if self.particles.x.unit == nbody_system.length else constants.G
return self.particles.potential_energy(G=G)
@property
def kinetic_energy(self):
return self.particles.kinetic_energy()
class GravityCodeForTesting(object):
def __init__(self):
self.particles = Particles()
self.model_time = quantities.zero
def evolve_model(self, t_end):
self.particles.position += self.particles.velocity * (t_end - self.model_time)
self.model_time = t_end
@property
def potential_energy(self):
G = nbody_system.G if self.particles.x.unit == nbody_system.length else constants.G
return self.particles.potential_energy(G=G)
@property
def kinetic_energy(self):
return self.particles.kinetic_energy()
class ExampleGravityCodeInterface(object):
def __init__(self, softening_mode="shared"):
self.particles = Particles()
if softening_mode == "individual":
self.softening_mode = "individual"
self._softening_lengths_squared = self._softening_lengths_squared_individual
self._softening_lengths = self._softening_lengths_individual
else:
self.softening_mode = "shared"
self._softening_lengths_squared = self._softening_lengths_squared_shared
self._softening_lengths = self._softening_lengths_shared
epsilon_squared_parameter = parameters.ModuleMethodParameterDefinition(
"get_epsilon_squared",
"set_epsilon_squared",
"epsilon_squared",
"gravitational softening length squared",
default_value = 0.0 | nbody_system.length**2,
must_set_before_get = False
)
self.parameters = parameters.new_parameters_instance_with_docs([epsilon_squared_parameter], self)
self.epsilon_squared = 0.0 | nbody_system.length**2
def _softening_lengths_squared_individual(self):
return self.particles.radius**2
def _softening_lengths_squared_shared(self):
return self.epsilon_squared.as_vector_with_length(len(self.particles))
def _softening_lengths_individual(self):
return self.particles.radius
def _softening_lengths_shared(self):
return self.epsilon_squared.sqrt().as_vector_with_length(len(self.particles))
def initialize_code(self):
self.model_time = 0 | units.Myr
def get_potential_at_point(self, eps, x ,y, z):
if isinstance(x, VectorQuantity):
return -constants.G * (self.particles.mass.reshape((-1,1)) /
(self._softening_lengths_squared().reshape((-1,1)) + eps.reshape((1,-1))**2 + (self.particles.x.reshape((-1,1)) -
x.reshape((1,-1)))**2 + (self.particles.y.reshape((-1,1)) - y.reshape((1,-1)))**2 +
(self.particles.z.reshape((-1,1)) - z.reshape((1,-1)))**2).sqrt()).sum(axis=0)
return -constants.G * (self.particles.mass / (self._softening_lengths_squared() + eps**2 + (self.particles.x - x)**2 +
(self.particles.y - y)**2 + (self.particles.z - z)**2).sqrt()).sum()
def get_gravity_at_point(self, eps, x ,y, z):
if isinstance(x, VectorQuantity):
delta_x = x.reshape((1,-1)) - self.particles.x.reshape((-1,1))
delta_y = y.reshape((1,-1)) - self.particles.y.reshape((-1,1))
delta_z = z.reshape((1,-1)) - self.particles.z.reshape((-1,1))
factor = -constants.G * (self.particles.mass.reshape((-1,1)) /
(self._softening_lengths_squared().reshape((-1,1)) + eps.reshape((1,-1))**2 +
delta_x**2 + delta_y**2 + delta_z**2)**1.5)
return (factor*delta_x).sum(axis=0), (factor*delta_y).sum(axis=0), (factor*delta_z).sum(axis=0)
delta_x = self.particles.x - x
delta_y = self.particles.y - y
delta_z = self.particles.z - z
factor = -constants.G * (self.particles.mass /
(self._softening_lengths_squared() + eps**2 + delta_x**2 + delta_y**2 + delta_z**2)**1.5)
return (factor*delta_x).sum(), (factor*delta_y).sum(), (factor*delta_z).sum()
@property
def potential_energy(self):
if self.softening_mode == "individual":
if len(self.particles) < 2:
return zero * constants.G
eps_vector = self.particles.radius**2
sum_of_energies = zero
for i in range(len(self.particles) - 1):
dx = self.particles[i].x - self.particles[i+1:].x
dy = self.particles[i].y - self.particles[i+1:].y
dz = self.particles[i].z - self.particles[i+1:].z
dr = ((dx * dx) + (dy * dy) + (dz * dz) + eps_vector[i] + eps_vector[i+1:]).sqrt()
sum_of_energies -= (self.particles.mass[i] * self.particles.mass[i+1:] / dr).sum()
return constants.G * sum_of_energies
return self.particles.potential_energy(smoothing_length_squared=2*self.epsilon_squared)
def before_set_parameter(self):
pass
def before_get_parameter(self):
pass
@property
def kinetic_energy(self):
return self.particles.kinetic_energy()
def get_epsilon_squared(self):
return self.epsilon_squared
def set_epsilon_squared(self, epsilon_squared):
self.epsilon_squared = epsilon_squared
def commit_particles(self):
self.set_accelerations()
self.set_next_timestep()
def set_accelerations(self):
accelerations = self.get_gravity_at_point(self._softening_lengths(),
self.particles.x, self.particles.y, self.particles.z)
self.particles.ax = accelerations[0]
self.particles.ay = accelerations[1]
self.particles.az = accelerations[2]
def set_next_timestep(self):
self.next_timestep = min([0.01 * (self.particles.velocity /
self.particles.acceleration).lengths_squared().amin().sqrt(), 1 | units.yr])
def evolve_model(self, t_end):
while self.model_time < t_end:
dt = self.next_timestep
self.particles.position += self.particles.velocity * dt + 0.5 * self.particles.acceleration * dt**2
old_acceleration = self.particles.acceleration
self.set_accelerations()
self.particles.velocity += 0.5 * (old_acceleration + self.particles.acceleration) * dt
self.model_time += dt
self.set_next_timestep()
class ExampleGravityCodeInterface(object):
def __init__(self, softening_mode="shared"):
self.particles = Particles()
if softening_mode == "individual":
self.softening_mode = "individual"
self._softening_lengths_squared = self._softening_lengths_squared_individual
self._softening_lengths = self._softening_lengths_individual
else:
self.softening_mode = "shared"
self._softening_lengths_squared = self._softening_lengths_squared_shared
self._softening_lengths = self._softening_lengths_shared
epsilon_squared_parameter = parameters.ModuleMethodParameterDefinition(
"get_epsilon_squared",
"set_epsilon_squared",
"epsilon_squared",
"gravitational softening length squared",
default_value=0.0 | nbody_system.length ** 2,
must_set_before_get=False,
)
self.parameters = parameters.new_parameters_instance_with_docs([epsilon_squared_parameter], self)
self.epsilon_squared = 0.0 | nbody_system.length ** 2
def _softening_lengths_squared_individual(self):
return self.particles.radius ** 2
def _softening_lengths_squared_shared(self):
return self.epsilon_squared.as_vector_with_length(len(self.particles))
def _softening_lengths_individual(self):
return self.particles.radius
def _softening_lengths_shared(self):
return self.epsilon_squared.sqrt().as_vector_with_length(len(self.particles))
def initialize_code(self):
self.model_time = 0 | units.Myr
def get_potential_at_point(self, eps, x, y, z):
if isinstance(x, VectorQuantity):
return -constants.G * (
self.particles.mass.reshape((-1, 1))
/ (
self._softening_lengths_squared().reshape((-1, 1))
+ eps.reshape((1, -1)) ** 2
+ (self.particles.x.reshape((-1, 1)) - x.reshape((1, -1))) ** 2
+ (self.particles.y.reshape((-1, 1)) - y.reshape((1, -1))) ** 2
+ (self.particles.z.reshape((-1, 1)) - z.reshape((1, -1))) ** 2
).sqrt()
).sum(axis=0)
return (
-constants.G
* (
self.particles.mass
/ (
self._softening_lengths_squared()
+ eps ** 2
+ (self.particles.x - x) ** 2
+ (self.particles.y - y) ** 2
+ (self.particles.z - z) ** 2
).sqrt()
).sum()
)
def get_gravity_at_point(self, eps, x, y, z):
if isinstance(x, VectorQuantity):
delta_x = x.reshape((1, -1)) - self.particles.x.reshape((-1, 1))
delta_y = y.reshape((1, -1)) - self.particles.y.reshape((-1, 1))
delta_z = z.reshape((1, -1)) - self.particles.z.reshape((-1, 1))
factor = -constants.G * (
self.particles.mass.reshape((-1, 1))
/ (
self._softening_lengths_squared().reshape((-1, 1))
+ eps.reshape((1, -1)) ** 2
+ delta_x ** 2
+ delta_y ** 2
+ delta_z ** 2
)
** 1.5
)
return (factor * delta_x).sum(axis=0), (factor * delta_y).sum(axis=0), (factor * delta_z).sum(axis=0)
delta_x = self.particles.x - x
delta_y = self.particles.y - y
delta_z = self.particles.z - z
factor = -constants.G * (
self.particles.mass
/ (self._softening_lengths_squared() + eps ** 2 + delta_x ** 2 + delta_y ** 2 + delta_z ** 2) ** 1.5
)
return (factor * delta_x).sum(), (factor * delta_y).sum(), (factor * delta_z).sum()
@property
def potential_energy(self):
if self.softening_mode == "individual":
if len(self.particles) < 2:
return zero * constants.G
eps_vector = self.particles.radius ** 2
sum_of_energies = zero
for i in range(len(self.particles) - 1):
dx = self.particles[i].x - self.particles[i + 1 :].x
dy = self.particles[i].y - self.particles[i + 1 :].y
dz = self.particles[i].z - self.particles[i + 1 :].z
dr = ((dx * dx) + (dy * dy) + (dz * dz) + eps_vector[i] + eps_vector[i + 1 :]).sqrt()
sum_of_energies -= (self.particles.mass[i] * self.particles.mass[i + 1 :] / dr).sum()
return constants.G * sum_of_energies
return self.particles.potential_energy(smoothing_length_squared=2 * self.epsilon_squared)
def before_set_parameter(self):
pass
def before_get_parameter(self):
pass
@property
def kinetic_energy(self):
return self.particles.kinetic_energy()
def get_epsilon_squared(self):
return self.epsilon_squared
def set_epsilon_squared(self, epsilon_squared):
self.epsilon_squared = epsilon_squared
def commit_particles(self):
self.set_accelerations()
self.set_next_timestep()
def set_accelerations(self):
accelerations = self.get_gravity_at_point(
self._softening_lengths(), self.particles.x, self.particles.y, self.particles.z
)
self.particles.ax = accelerations[0]
self.particles.ay = accelerations[1]
self.particles.az = accelerations[2]
def set_next_timestep(self):
self.next_timestep = min(
[
0.01 * (self.particles.velocity / self.particles.acceleration).lengths_squared().amin().sqrt(),
1 | units.yr,
]
)
def evolve_model(self, t_end):
while self.model_time < t_end:
dt = self.next_timestep
self.particles.position += self.particles.velocity * dt + 0.5 * self.particles.acceleration * dt ** 2
old_acceleration = self.particles.acceleration
self.set_accelerations()
self.particles.velocity += 0.5 * (old_acceleration + self.particles.acceleration) * dt
self.model_time += dt
self.set_next_timestep()
