"""
In this script we simulate a plummer sphere on a grid
"""
from amuse.support.core import late
from amuse.units import constants
from amuse.units import units
from amuse.units import nbody_system
from amuse.units.quantities import VectorQuantity
from amuse.units.generic_unit_system import *
from amuse.units.generic_unit_converter import *
from amuse.ext import cloud
from amuse.community.athena.interface import Athena, AthenaInterface
from amuse.community.capreole.interface import Capreole
from amuse.community.hermite0.interface import Hermite
from amuse.community.phiGRAPE.interface import PhiGRAPE
from amuse.community.octgrav.interface import Octgrav
import sys
try:
from amuse import plot
from matplotlib import pyplot
IS_PLOT_AVAILABLE = True
except ImportError:
IS_PLOT_AVAILABLE = False
from numpy import sqrt, arange, searchsorted, tanh, pi
from optparse import OptionParser
from amuse import io
from amuse.io import text
from amuse.datamodel import Grid
from amuse.datamodel import Particles
class HydroGridAndNbody(object):
def __init__(self, gridcode, nbodycode):
self.gridcode = gridcode
self.nbodycode = nbodycode
def setup_positions_for_potential_grid(self):
self.x = self.gridcode.potential_grid.x.flatten()
self.y = self.gridcode.potential_grid.y.flatten()
self.z = self.gridcode.potential_grid.z.flatten()
self.eps = self.x.aszeros()
def setup_particles_in_nbodycode(self):
staggered_grid_shape = numpy.asarray(self.gridcode.grid.shape) + 2
corner0 = self.gridcode.grid[0][0][0].position
corner1 = self.gridcode.grid[-1][-1][-1].position
delta = self.gridcode.grid[1][1][1].position - corner0
print delta.prod()
self.volume = delta.prod()
staggered_corner0 = corner0 - delta
staggered_corner1 = corner1
self.staggered_grid = Grid.create(staggered_grid_shape, staggered_corner1 - staggered_corner0 + delta)
#self.staggered_grid.x += staggered_corner0.x
#self.staggered_grid.y += staggered_corner0.x
#self.staggered_grid.z += staggered_corner0.x
#self.staggered_grid.p000 = 0.0 | mass
#self.staggered_grid.p100 = 0.0 | mass
#self.staggered_grid.p010 = 0.0 | mass
#self.staggered_grid.p110 = 0.0 | mass
#self.staggered_grid.p000 = 0.0 | mass
#self.staggered_grid.p100 = 0.0 | mass
#self.staggered_grid.p011 = 0.0 | mass
#self.staggered_grid.p111 = 0.0 | mass
#self.staggered_grid.p001 = 0.0 | mass
#self.staggered_grid.p101 = 0.0 | mass
self.normal_grid = Grid.create(self.gridcode.grid.shape, corner1 - corner0 + delta)
particles = Particles(self.normal_grid.size)
particles.mass = 0.0 | mass
particles.x = self.grid.x.flatten()
particles.y = self.grid.y.flatten()
particles.z = self.grid.z.flatten()
particles.vx = 0.0 | speed
particles.vy = 0.0 | speed
particles.vz = 0.0 | speed
particles.radius = 0.0 | length
self.nbodycode.particles.add_particles(particles)
self.from_model_to_nbody = particles.new_channel_to(self.nbodycode.particles)
self.nbodycode.commit_particles()
self.particles = particles
@property
def grid(self):
return self.gridcode.grid
@property
def parameters(self):
return self.gridcode.parameters
def commit_parameters(self):
self.gridcode.commit_parameters()
self.setup_positions_for_potential_grid()
self.setup_particles_in_nbodycode()
def initialize_grid(self):
self.gridcode.initialize_grid()
self.position = self.gridcode.grid.position
def evolve_model(self, time):
#masses = self.gridcode.grid.rho * self.volume * 1.0 / 8.0
#self.staggered_grid.p000[1:,1:,1:] = masses
#self.staggered_grid.p100[:-1,1:,1:] = masses
# self.staggered_grid.p010[1:,:-1,1:] = masses
#self.staggered_grid.p110[:-1,:-1,1:] = masses
#self.staggered_grid.p001[1:,1:,:-1] = masses
#self.staggered_grid.p101[:-1,1:,:-1] = masses
#self.staggered_grid.p011[1:,:-1,:-1] = masses
#self.staggered_grid.p111[:-1,:-1,:-1] = masses
#self.staggered_grid.masses = (
# self.staggered_grid.p000 +
# self.staggered_grid.p100 +
# self.staggered_grid.p010 +
# self.staggered_grid.p110 +
# self.staggered_grid.p001 +
# self.staggered_grid.p101 +
# self.staggered_grid.p011 +
# self.staggered_grid.p111
#)
#self.particles.mass = self.staggered_grid.masses.flatten()
self.particles.mass = (self.gridcode.grid.rho * self.volume).flatten()
self.from_model_to_nbody.copy_attribute('mass')
correction = (length ** 3) / (mass * (time ** 2))
print correction, nbody_system.G
print "getting potential energy"
potential = self.nbodycode.get_potential_at_point(
self.eps,
self.x,
self.y,
self.z
)
print potential.shape, self.gridcode.potential_grid.shape, self.grid.shape
self.staggered_grid.rho = 0.0 | density
self.staggered_grid[1:-1,1:-1,1:-1].rho = self.gridcode.grid.rho
correction = ((self.staggered_grid.rho * self.volume) / numpy.sqrt(self.nbodycode.parameters.epsilon_squared)) * nbody_system.G
print correction.flatten().shape, potential.shape
potential = potential + correction.flatten()
print "got potential enery"
potential = potential.reshape(self.gridcode.potential_grid.shape)
self.gridcode.potential_grid.potential = potential
print "evolve hydro", time
self.gridcode.evolve_model(time)
print "end evolve hydro"
class HydroGridAndNbodyWithAccelerationTransfer(object):
def __init__(self, gridcode, nbodycode):
self.gridcode = gridcode
self.nbodycode = nbodycode
self.gridcode.grid.add_vector_attribute("acceleration_field", ["fx","fy","fz"])
def setup_positions_for_potential_grid(self):
self.x = self.gridcode.acceleration_grid.x.flatten()
self.y = self.gridcode.acceleration_grid.y.flatten()
self.z = self.gridcode.acceleration_grid.z.flatten()
self.eps = self.x.aszeros()
def setup_particles_in_nbodycode(self):
staggered_grid_shape = numpy.asarray(self.gridcode.grid.shape) + 1
corner0 = self.gridcode.grid[0][0][0].position
corner1 = self.gridcode.grid[-1][-1][-1].position
delta = self.gridcode.grid[1][1][1].position - corner0
print delta.prod()
self.volume = delta.prod()
staggered_corner0 = corner0 - delta
staggered_corner1 = corner1
self.staggered_grid = Grid.create(staggered_grid_shape, staggered_corner1 - staggered_corner0 + delta)
self.staggered_grid.x += staggered_corner0.x
self.staggered_grid.y += staggered_corner0.x
self.staggered_grid.z += staggered_corner0.x
self.staggered_grid.p000 = 0.0 | mass
self.staggered_grid.p100 = 0.0 | mass
self.staggered_grid.p010 = 0.0 | mass
self.staggered_grid.p110 = 0.0 | mass
self.staggered_grid.p000 = 0.0 | mass
self.staggered_grid.p100 = 0.0 | mass
self.staggered_grid.p011 = 0.0 | mass
self.staggered_grid.p111 = 0.0 | mass
self.staggered_grid.p001 = 0.0 | mass
self.staggered_grid.p101 = 0.0 | mass
self.normal_grid = Grid.create(self.gridcode.grid.shape, corner1 - corner0 + delta)
particles = Particles(self.staggered_grid.size)
particles.mass = 0.0 | mass
particles.x = self.normal_grid.x.flatten()
particles.y = self.normal_grid.y.flatten()
particles.z = self.normal_grid.z.flatten()
particles.vx = 0.0 | speed
particles.vy = 0.0 | speed
particles.vz = 0.0 | speed
particles.radius = 0.0 | length
self.nbodycode.particles.add_particles(particles)
self.from_model_to_nbody = particles.new_channel_to(self.nbodycode.particles)
self.nbodycode.commit_particles()
self.particles = particles
@property
def grid(self):
return self.gridcode.grid
@property
def parameters(self):
return self.gridcode.parameters
def commit_parameters(self):
self.gridcode.commit_parameters()
self.setup_positions_for_potential_grid()
self.setup_particles_in_nbodycode()
def initialize_grid(self):
self.gridcode.initialize_grid()
self.position = self.gridcode.grid.position
def evolve_model(self, time):
masses = self.gridcode.grid.rho * self.volume * 1.0 / 8.0
self.staggered_grid.p000[1:,1:,1:] = masses
self.staggered_grid.p100[:-1,1:,1:] = masses
self.staggered_grid.p010[1:,:-1,1:] = masses
self.staggered_grid.p110[:-1,:-1,1:] = masses
self.staggered_grid.p001[1:,1:,:-1] = masses
self.staggered_grid.p101[:-1,1:,:-1] = masses
self.staggered_grid.p011[1:,:-1,:-1] = masses
self.staggered_grid.p111[:-1,:-1,:-1] = masses
self.staggered_grid.masses = (
self.staggered_grid.p000 +
self.staggered_grid.p100 +
self.staggered_grid.p010 +
self.staggered_grid.p110 +
self.staggered_grid.p001 +
self.staggered_grid.p101 +
self.staggered_grid.p011 +
self.staggered_grid.p111
)
#self.particles.mass = self.staggered_grid.masses.flatten()
self.particles.mass = self.gridcode.grid.rho * self.volume
self.from_model_to_nbody.copy_attribute('mass')
print "getting acceleration field"
acc_x, acc_y, acc_z = self.nbodycode.get_gravity_at_point(
self.eps,
self.x,
self.y,
self.z
)
print "got acceleration field"
acc_x = acc_x.reshape(self.grid.shape)
acc_y = acc_y.reshape(self.grid.shape)
acc_z = acc_z.reshape(self.grid.shape)
self.gridcode.acceleration_grid.set_values_in_store(None, ["fx","fy","fz"], [acc_x, acc_y, acc_z])
print "evolve hydro", time
self.gridcode.evolve_model(time)
print "end evolve hydro"
class CalculateSolutionIn3D(object):
size = 10.0 | length
number_of_workers = 1
number_of_grid_points = 25
gamma = 5.0/3.0
rho_medium = 0 | density
rho_sphere = 0.5 | density
center = [1.0 / 2.0, 1.0 / 2.0, 1.0 / 2.0] * size
radius = 1.0 | length
total_mass = 1.0 | mass
name_of_the_code = "athena_selfgravity"
convert_generic_units = ConvertBetweenGenericAndSiUnits(1.0 | units.kpc, 1.0e10 | units.MSun, constants.G)
def __init__(self, **keyword_arguments):
for x in keyword_arguments:
print x, keyword_arguments[x]
setattr(self, x, keyword_arguments[x])
self.dimensions_of_mesh = (
self.number_of_grid_points,
self.number_of_grid_points,
self.number_of_grid_points
)
self.instance = None
def new_instance_of_code(self):
attribute = "new_instance_of_{0}_code".format(self.name_of_the_code.lower())
return getattr(self,attribute)()
def new_instance_of_athena_selfgravity_code(self):
result=Athena(mode = AthenaInterface.MODE_SELF_GRAVITY, redirection="none", number_of_workers=1)
result.initialize_code()
result.parameters.gamma = self.gamma
result.parameters.courant_number=0.3
result.set_grav_mean_rho(self.rho_mean.value_in(density))
return result
def new_instance_of_athena_code(self):
result=Athena(redirection="none", number_of_workers=1)
result.initialize_code()
result.parameters.gamma = self.gamma
result.parameters.courant_number=0.3
return result
def new_instance_of_athena_and_nbody_code(self):
gridcode=Athena(redirection="none", number_of_workers=1)
gridcode.initialize_code()
gridcode.parameters.gamma = self.gamma
gridcode.parameters.courant_number=0.3
gridcode.set_has_external_gravitational_potential(1)
nbodycode = PhiGRAPE(mode="gpu")
#nbodycode = Octgrav(mode="gpu")
nbodycode.initialize_code()
nbodycode.parameters.epsilon_squared = (self.size / (10000.0 * self.number_of_grid_points)) ** 2
nbodycode.commit_parameters()
result = HydroGridAndNbody(gridcode, nbodycode)
return result
def new_instance_of_capreole_and_nbody_code(self):
gridcode=Capreole(number_of_workers=self.number_of_workers)
gridcode.initialize_code()
nbodycode = PhiGRAPE(mode="gpu")
#nbodycode = Octgrav(mode="gpu")
nbodycode.initialize_code()
nbodycode.parameters.epsilon_squared = (self.size / (10000.0 * self.number_of_grid_points)) ** 2
nbodycode.commit_parameters()
result = HydroGridAndNbodyWithAccelerationTransfer(gridcode, nbodycode)
return result
def set_parameters(self, instance):
instance.parameters.mesh_size = self.dimensions_of_mesh
instance.parameters.length_x = self.size
instance.parameters.length_y = self.size
instance.parameters.length_z = self.size
instance.parameters.x_boundary_conditions = ("periodic","periodic")
instance.parameters.y_boundary_conditions = ("periodic","periodic")
instance.parameters.z_boundary_conditions = ("periodic","periodic")
result = instance.commit_parameters()
def new_grid(self):
density = mass / length**3
density = density
momentum = speed * density
energy = mass / (time**2 * length)
grid = Grid.create(self.dimensions_of_mesh, (10.0, 10.0, 10.0) | length)
grid.rho = self.rho_medium
grid.rhovx = 0.0 | momentum
grid.rhovy = 0.0 | momentum
grid.rhovz = 0.0 | momentum
grid.energy = 0.0 | energy
return grid
def initialize_grid_with_plummer_sphere(self, grid):
scaled_radius = self.radius #/ 1.695
radii = (grid.position - self.center).lengths()
selected_radii = radii[radii < self.radius ]
print "number of cells in cloud (number of cells in grid)", len(selected_radii), grid.shape, grid.size
self.rho_sphere = ((0.75 * self.total_mass / (pi * (scaled_radius ** 3))))
grid.rho = self.rho_medium + (self.rho_sphere * ((1 + (radii ** 2) / (scaled_radius ** 2))**(-5.0/2.0)))
internal_energy = (0.25 | potential) * ( 1.0 | length / mass) * self.total_mass / scaled_radius
# 1.0 * grid.rho *
grid.energy = grid.rho * (internal_energy/((1.0+(radii/scaled_radius)**2)**(1.0/2.0)))
def setup_code(self):
print "setup code"
self.grid = self.new_grid()
self.initialize_grid_with_plummer_sphere(self.grid)
print "Mean density", self.grid.rho.flatten().mean()
self.rho_mean = self.grid.rho.flatten().mean()
self.instance=self.new_instance_of_code()
self.set_parameters(self.instance)
self.from_model_to_code = self.grid.new_channel_to(self.instance.grid)
self.from_code_to_model = self.instance.grid.new_channel_to(self.grid)
self.from_code_to_model.copy_attributes(("x","y","z",))
self.from_model_to_code.copy()
self.instance.initialize_grid()
def get_solution_at_time(self, time):
if self.instance is None:
self.setup_code()
print "start evolve"
self.instance.evolve_model(time)
print "copying results"
self.from_code_to_model.copy()
#print "max,min", max(self.grid.rhovx.flatten()), min(self.grid.rhovx.flatten())
#print "max,min", max(self.grid.rhovy.flatten()), min(self.grid.rhovy.flatten())
#print "max,min", max(self.grid.rhovz.flatten()), min(self.grid.rhovz.flatten())
#print "max,min", max(self.grid.energy.flatten()), min(self.grid.energy.flatten())
#print "max,min", max(self.grid.rho.flatten()), min(self.grid.rho.flatten())
return self.grid
def store_attributes(x, rho, rhovx, energy, filename):
output = text.TableFormattedText(filename = filename)
output.quantities = (x, rho, rhovx, energy)
output.attribute_names = ("x", "rho", "rhovx", "energy")
output.store()
def store_attributes_of_line(grid, yindex = 0, zindex = 0, **options):
store_attributes(
grid.x[...,yindex,zindex],
grid.rho[...,yindex,zindex],
grid.rhovx[...,yindex,zindex],
grid.energy[...,yindex,zindex],
filename = "plummer_sphere_{name_of_the_code}_{number_of_grid_points}_{number_of_workers}_{time}.csv".format(**options)
)
def new_option_parser():
result = OptionParser()
result.add_option(
"-n",
"--mesh-size",
dest="number_of_grid_points",
type="int",
default = 25,
help="number of grid cells in the x, y and z direction"
)
result.add_option(
"-w",
"--workers",
dest="number_of_workers",
type="int",
default = 1,
help="number of parallel workers to start"
)
result.add_option(
"-c",
"--code",
dest="name_of_the_code",
default="athena_selfgravity",
help="name of the code to use"
)
return result
def main(**options):
center = options['number_of_grid_points'] / 2
print "calculating shock using code"
model = CalculateSolutionIn3D(**options)
for t in range(50):
grid = model.get_solution_at_time( t * 0.1 | time)
print "saving data"
store_attributes_of_line(grid, yindex = center, zindex = center, time = t, **options)
if __name__ == "__main__":
options, arguments = new_option_parser().parse_args()
main(**options.__dict__)
Keywords: python, amuse, astrophysics, matplotlib, pylab, example, codex (see how-to-search-examples)