def __init__(self, particles = None, gravity_constant = None,
softening_mode="shared", G = None):
if particles is None:
self.particles=datamodel.Particles()
else:
self.particles = particles
if gravity_constant is None:
gravity_constant = G
elif not G is None:
raise Exception("both the parameter 'gravity_constant'({0}) and the parameter 'G'({1}) are given, please specify only one!".format(gravity_constant, G))
if gravity_constant is None:
if len(particles) and hasattr(particles, 'mass'):
try:
particles[0].mass.value_in(units.kg)
self.gravity_constant = constants.G
except:
raise AmuseException("For generic units the gravity_constant must be specified")
else:
raise AmuseException("Particle data not yet available, so the gravity_constant must be specified")
else:
self.gravity_constant = gravity_constant
if softening_mode == "individual" or softening_mode == "radius":
self._softening_lengths_squared = self._softening_lengths_squared_individual
elif softening_mode == "h_smooth":
self._softening_lengths_squared = self._softening_lengths_squared_h_smooth
else:
self._softening_lengths_squared = self._softening_lengths_squared_shared
self.smoothing_length_squared = quantities.zero
def __init__(self,
collision_code,
collision_code_arguments=dict(),
collision_code_parameters=dict(),
gravity_code=None,
stellar_evolution_code=None,
verbose=False,
**options):
if inspect.isclass(collision_code):
self.collision_code_name = collision_code.__name__
else:
self.collision_code_name = collision_code.__class__.__name__
if collision_code.stellar_evolution_code_required and stellar_evolution_code is None:
raise AmuseException(
"{0} requires a stellar evolution code: "
"CollisionHandler(..., stellar_evolution_code=x)".format(
self.collision_code_name))
if collision_code.gravity_code_required and gravity_code is None:
raise AmuseException(
"{0} requires a gravity code: "
"CollisionHandler(..., gravity_code=x)".format(
self.collision_code_name))
self.collision_code = collision_code
self.collision_code_arguments = collision_code_arguments
self.collision_code_parameters = collision_code_parameters
self.gravity_code = gravity_code
self.stellar_evolution_code = stellar_evolution_code
self.verbose = verbose
self.options = options
def setup_core_parameters(self):
if self.target_core_mass is None:
if hasattr(self.particle, "core_mass"):
self.target_core_mass = self.particle.core_mass
else:
raise AmuseException(
"Requested model has with_core_particle=True, but no target_core_mass specified."
)
self.original_entropy = (
self.gamma - 1.0) * (self.specific_internal_energy_profile *
self.density_profile**(1.0 - self.gamma))
interpolator = EnclosedMassInterpolator()
interpolator.initialize(self.radius_profile, self.density_profile)
i_edge = numpy.searchsorted(
interpolator.enclosed_mass.number,
self.target_core_mass.value_in(
interpolator.enclosed_mass.unit)) - 1
min_i = i_edge
max_i = len(self.radius_profile) - 3
enclosed_mass_edge = interpolator.enclosed_mass[min_i + 1]
min_enclosed_mass_residual = self.construct_model_with_core(
min_i, enclosed_mass_edge, self.gamma)
enclosed_mass_edge = interpolator.enclosed_mass[max_i + 1]
max_enclosed_mass_residual = self.construct_model_with_core(
max_i, enclosed_mass_edge, self.gamma)
while (max_enclosed_mass_residual < zero) and (max_i > min_i + 20):
max_i -= 20
enclosed_mass_edge = interpolator.enclosed_mass[max_i + 1]
max_enclosed_mass_residual = self.construct_model_with_core(
max_i, enclosed_mass_edge, self.gamma)
if (min_enclosed_mass_residual > zero) or (max_enclosed_mass_residual <
zero):
raise AmuseException(
"Requested target_core_mass of {0} is out of range.".format(
self.target_core_mass))
while max_i - min_i > 1:
next_i = (max_i + min_i) / 2
enclosed_mass_edge = interpolator.enclosed_mass[next_i + 1]
enclosed_mass_residual = self.construct_model_with_core(
next_i, enclosed_mass_edge, self.gamma)
if enclosed_mass_residual >= zero:
max_i = next_i
else:
min_i = next_i
if enclosed_mass_residual < zero:
enclosed_mass_edge = interpolator.enclosed_mass[max_i + 1]
self.construct_model_with_core(max_i, enclosed_mass_edge,
self.gamma)
self.density_profile = self.rho
self.specific_internal_energy_profile = self.u
interpolator.initialize(self.radius_profile, self.density_profile)
self.core_mass = self.mass - interpolator.enclosed_mass[-1]
self.core_radius = self.radius_profile[max_i] / 2.8
self.mass = self.mass - self.core_mass
def new_particles_with_blackbody_color(original_particles, **kwargs):
"""
Returns new color particles. These are bound to the 'original_particles' in
the sense that they share their attributes, but have additional attributes
'red', 'green', and 'blue'. These colors are based on Mitchell Charity's
blackbody color datafile (bbr_color.txt, see
http://www.vendian.org/mncharity/dir3/blackbody/)
If the particles have a temperature attribute, the colors are computed from
these. Otherwise they will be computed from the gas internal energy, using the
T_from_u function, in which case the optional keyword arguments X, Y,
Z, and x_ion can be supplied.
:argument original_particles: the particles for which the color needs to be computed
:argument X: hydrogen abundance for T_from_u converter (default: None, i.e. compute from Y and Z)
:argument Y: helium abundance for T_from_u converter (default: 0.25)
:argument Z: metal (everything heavier than helium) abundance for T_from_u converter (default: 0.02)
:argument x_ion: ionisation fraction for T_from_u converter (default: 0.1)
"""
if "temperature" in original_particles.get_attribute_names_defined_in_store(
):
colors = BlackBodyColorFromTemperature(**kwargs)
attributes_names = ["temperature"]
elif "u" in original_particles.get_attribute_names_defined_in_store():
colors = BlackBodyColorFromInternalEnergy(**kwargs)
attributes_names = ["u"]
else:
raise AmuseException(
"The particles need to have 'temperature' or 'u' attributes for deriving black body colors"
)
return new_particles_with_color(original_particles,
colors.red_function,
colors.green_function,
colors.blue_function,
attributes_names=attributes_names)
def convert_comoving_to_physical(original, redshift=0.0, hubble_parameter=1.0, attribute_names=None):
"""
Converts quantities or particle sets from comoving coordinates to physical
coordinates. In comoving coordinates, changes in positions due to the
expansion of the universe are corrected for, by dividing by the scale
factor 'a':
a = 1 / (1 + z).
This function will undo this correction for all quantities with units that
are (derived from) length units (e.g. position, velocity, energy, but not
time or mass).
Optionally, a value for the Hubble parameter (value of Hubble constant in
units of 100 km/s/Mpc) can be supplied. If so, the units of 'original' are
assumed to be based on length/h, mass/h, and time/h, instead of length,
mass, and time, respectively. These factors will be divided out in the result
"""
if isinstance(original, Particles):
copy = original.copy()
if attribute_names is None:
attribute_names = copy.get_attribute_names_defined_in_store()
for attribute in attribute_names:
setattr(copy, attribute, convert_quantity_from_comoving_to_physical(
getattr(copy, attribute), redshift, hubble_parameter))
return copy
elif hasattr(original, "unit"):
return convert_quantity_from_comoving_to_physical(original, redshift, hubble_parameter)
else:
raise AmuseException("Can't convert instance of {0} from comoving to physical "
"coordinates (only Particles or Quantity supported)".format(original.__class__))
def new_stellar_wind(sph_particle_mass, target_gas=None, timestep=None,
derive_from_evolution=False, mode="simple", **kwargs):
"""
Create a new stellar wind code.
target_gas: a gas particle set into which the wind particles should be
put (requires timestep)
timestep: the timestep at which the wind particles should be generated.
derive_from_evolution: derive the wind parameters from stellar
evolution (you still need to update the stellar parameters)
mode: set to 'simple', 'accelerate' or 'mechanical'
"""
if (target_gas is None) ^ (timestep is None):
raise AmuseException("Must specify both target_gas and timestep"
"(or neither)")
wind_modes = {"simple": SimpleWind,
"accelerate": AcceleratingWind,
"heating": HeatingWind,
}
stellar_wind = wind_modes[mode](sph_particle_mass, derive_from_evolution,
**kwargs)
if target_gas is not None:
stellar_wind.set_target_gas(target_gas, timestep)
return stellar_wind
def convert_particles_to_pynbody_data(particles, length_unit=units.kpc, pynbody_unit="kpc"):
if not HAS_PYNBODY:
raise AmuseException("Couldn't find pynbody")
if hasattr(particles, "u"):
pynbody_data = new(gas=len(particles))
else:
pynbody_data = new(dm=len(particles))
pynbody_data._filename = "AMUSE"
if hasattr(particles, "mass"):
pynbody_data['mass'] = SimArray(particles.mass.value_in(units.MSun), "Msol")
if hasattr(particles, "position"):
pynbody_data['x'] = SimArray(particles.x.value_in(length_unit), pynbody_unit)
pynbody_data['y'] = SimArray(particles.y.value_in(length_unit), pynbody_unit)
pynbody_data['z'] = SimArray(particles.z.value_in(length_unit), pynbody_unit)
if hasattr(particles, "velocity"):
pynbody_data['vx'] = SimArray(particles.vx.value_in(units.km / units.s), "km s^-1")
pynbody_data['vy'] = SimArray(particles.vy.value_in(units.km / units.s), "km s^-1")
pynbody_data['vz'] = SimArray(particles.vz.value_in(units.km / units.s), "km s^-1")
if hasattr(particles, "h_smooth"):
pynbody_data['smooth'] = SimArray(particles.h_smooth.value_in(length_unit), pynbody_unit)
if hasattr(particles, "rho"):
pynbody_data['rho'] = SimArray(particles.rho.value_in(units.g / units.cm**3),
"g cm^-3")
if hasattr(particles, "temp"):
pynbody_data['temp'] = SimArray(particles.temp.value_in(units.K), "K")
elif hasattr(particles, "u"):
# pynbody_data['u'] = SimArray(particles.u.value_in(units.km**2 / units.s**2), "km^2 s^-2")
temp = 2.0/3.0 * particles.u * mu() / constants.kB
pynbody_data['temp'] = SimArray(temp.value_in(units.K), "K")
return pynbody_data
def create_initial_wind(self, number=None, time=None, check_length=True):
"""
This is a convenience method that creates some initial particles.
They are created as if the wind has already been blowing for
'time'. Note that this does not work if the mass loss is derived
from stellar evolution.
If 'number' is given, the required time to get that number of
particles is calculated. This assumes that the number of expected
particles is far larger then the number of stars
"""
if number is not None:
required_mass = number * self.sph_particle_mass
total_mass_loss = self.particles.wind_mass_loss_rate.sum()
time = 1.0 * required_mass / total_mass_loss
self.model_time = time
self.particles.evolve_mass_loss(self.model_time)
if self.has_new_wind_particles():
wind_gas = self.create_wind_particles()
if self.has_target():
self.target_gas.add_particles(wind_gas)
elif check_length:
raise AmuseException("create_initial_wind time was too small to"
"create any particles.")
else:
wind_gas = Particles()
self.reset()
return wind_gas
def convert_quantity(self, quantity):
if has_nbody_unit(quantity.unit):
if self.nbody_converter:
return _quantity_in_preferred_units(self.preferred_units, self.nbody_converter.to_si(quantity))
elif not self.ignore_converter_exceptions:
raise AmuseException("Unable to convert {0} to SI units. No "
"nbody_converter given".format(quantity.unit))
return _quantity_in_preferred_units(self.preferred_units, quantity)
def get_index(self, value, sorted_vector):
out_of_bounds = numpy.logical_or(sorted_vector[0] > value, value > sorted_vector[-1])
if out_of_bounds.any():
value = numpy.compress(numpy.array([out_of_bounds]).flatten(), value.number) | value.unit
raise AmuseException("Can't find a valid index. {0} is not in "
"the range [{1}, {2}].".format(value, sorted_vector[0], sorted_vector[-1]))
index = numpy.searchsorted(sorted_vector.number, value.value_in(sorted_vector.unit))
return numpy.maximum(index - 1, 0)
def get_index(self, value, sorted_vector):
if not sorted_vector[0] <= value <= sorted_vector[-1]:
raise AmuseException("Can't find a valid index. {0} is not in "
"the range [{1}, {2}].".format(
value, sorted_vector[0],
sorted_vector[-1]))
index = numpy.searchsorted(sorted_vector, value)
return max(index - 1, 0)
def convert_quantity(self, quantity):
if is_not_nbody_unit(quantity.unit):
if self.nbody_converter:
return self.nbody_converter.to_nbody(quantity)
elif not self.ignore_converter_exceptions:
raise AmuseException("Unable to convert {0} to N-body units. No "
"nbody_converter given".format(quantity.unit))
return quantity
def get_indices(self, values, sorted_vector):
values = numpy.array(values.value_in(sorted_vector.unit))
sorted_vector = sorted_vector.number
if values.min() < sorted_vector[0] or values.max() > sorted_vector[-1]:
raise AmuseException("Can't find a valid index. Value not in "
"the range [{0}, {1}].".format(sorted_vector[0], sorted_vector[-1]))
indices = numpy.maximum(numpy.searchsorted(sorted_vector, values) - 1, 0)
return indices
def get_radius_for_enclosed_mass(self, enclosed_mass):
if not self.initialized:
raise AmuseException(
"Can't calculate radius for enclosed mass: interpolator is not initialized"
)
index = self.get_index(enclosed_mass, self.enclosed_mass)
return (((enclosed_mass - self.enclosed_mass[index]) /
(self.four_thirds_pi * self.densities[index]) +
self.radii_cubed[index]))**(1.0 / 3.0)
def get_enclosed_mass(self, radius):
if not self.initialized:
raise AmuseException(
"Can't calculate enclosed mass: interpolator is not initialized"
)
index = self.get_index(radius, self.radii)
return (self.enclosed_mass[index] +
self.four_thirds_pi * self.densities[index] *
(radius**3 - self.radii_cubed[index]))
def select(self, position, sources):
normalized_shape = self.dimensions / self.dimensions.max()
max_r_squared = self.dimensions.max()**2
rel_position = (sources.position-position)
if self.orientation is not None:
#TODO rotate relative_position using orientation
raise AmuseException("spheroid orientation not implemented yet")
return (rel_position/normalized_shape).lengths_squared() < max_r_squared
def has_log_scaling(array):
diff = numpy.diff(array)
if numpy.all(diff - diff[0] < diff[0]/10.):
return False
logdiff = numpy.diff(numpy.log10(array))
if numpy.all(logdiff - logdiff[0] < logdiff[0]/10.):
return True
raise AmuseException("This method cannot be used for non regular arrays")
def __init__(self, sph_code, dimensions, do_scale = False):
if (sph_code.mode != sph_code.MODE_PERIODIC_BOUNDARIES):
raise AmuseException("Only periodic boundary conditions supported")
if len(dimensions) != 3:
raise AmuseException("Argument dimensions must contain exactly three numbers")
if isinstance(sph_code, Fi):
self.box_offset = sph_code.parameters.periodic_box_size / 2.0
elif isinstance(sph_code, Gadget2):
self.box_offset = sph_code.parameters.periodic_box_size * 0.0
else:
raise AmuseException("Unknown hydrodynamics code: {0} - don't know whether the "
"box runs from 0 to L or from -0.5 L to 0.5 L.".format(sph_code.__class__.__name__))
self.sph_code = sph_code
self.dimensions = dimensions
self.do_scale = do_scale
self.box_size = sph_code.parameters.periodic_box_size
def ylim(*args, **kwargs):
if len(UnitlessArgs.arg_units) is 0:
raise AmuseException("Cannot call ylim function before plotting")
args = UnitlessArgs.value_in_y_unit(*args)
for name in ("ymin", "ymax"):
if name in kwargs:
kwargs[name] = UnitlessArgs.value_in_y_unit(kwargs[name])
native_plot.ylim(*args, **kwargs)
def mu(X = None, Y = 0.25, Z = 0.02, x_ion = 0.1):
"""
Compute the mean molecular weight in kg (the average weight of particles in a gas)
X, Y, and Z are the mass fractions of Hydrogen, of Helium, and of metals, respectively.
x_ion is the ionisation fraction (0 < x_ion < 1), 1 means fully ionised
"""
if X is None:
X = 1.0 - Y - Z
elif abs(X + Y + Z - 1.0) > 1e-6:
raise AmuseException("Error in calculating mu: mass fractions do not sum to 1.0")
return constants.proton_mass / (X*(1.0+x_ion) + Y*(1.0+2.0*x_ion)/4.0 + Z*x_ion/2.0)
def disable(self):
if self.is_supported():
self.conditions.code.disable_stopping_condition(self.type)
else:
name = [
name for name, value in self.conditions.all_conditions()
if value is self
][0]
raise AmuseException(
"Can't disable stopping condition '{0}', since '{1}' does not support this condition."
.format(name,
type(self.conditions.code).__name__))
def new_spherical_particle_distribution(
number_of_particles,
radial_density_func=None, # not yet supported, specify radii and densities tables:
radii=None,
densities=None,
total_mass=None,
size=None, # if total_mass is not given, it will be deduced from size or max(radii)
**keyword_arguments
): # optional arguments for UniformSphericalDistribution
"""
Returns a Particles set with positions following a spherical
distribution. The radial density profile is determined from the
look-up table (radii, densities). Entries in the 'radii' table
are interpreted as the outer radius of the shell, with uniform
density as defined by the corresponding entry in the 'densities'
table:
rho(r) = densities[i], for ( radii[i-1] <= r <= radii[i] )
Only the positions and masses (equal-mass system) are set.
:argument number_of_particles: Number of particles in the resulting model
:argument radii: Table with radii for the radial density profile
:argument densities: Table with densities for the radial density profile
:argument total_mass: Total mass of the Particles set (optional, will be
deduced from size or max(radii) otherwise)
:argument size: Radius of the sphere enclosing the model (optional)
"""
if (radii is None) or (densities is None):
raise AmuseException(
"Using an arbitrary radial density function is not yet "
"supported. Radius and density tables must be passed instead.")
interpolator = EnclosedMassInterpolator()
interpolator.initialize(radii, densities)
if total_mass is None:
total_mass = interpolator.get_enclosed_mass(size or max(radii))
particles = Particles(number_of_particles)
particle_mass = total_mass * 1.0 / number_of_particles
particles.mass = particle_mass
x, y, z = UniformSphericalDistribution(number_of_particles,
**keyword_arguments).result
# Now scale the uniformly distributed particle positions to match the radial density profile
r_old = numpy.sqrt(x * x + y * y + z * z)
indices = numpy.argsort(r_old)
if r_old[indices[0]] == 0.0:
r_old[indices[0]] = 1.0
f_scale = interpolator.get_radius_for_enclosed_mass(
(numpy.arange(0.5, number_of_particles + 0.5) | units.none) *
particle_mass) / r_old[indices]
particles.x = (f_scale * x[indices]).as_quantity_in(radii.unit)
particles.y = (f_scale * y[indices]).as_quantity_in(radii.unit)
particles.z = (f_scale * z[indices]).as_quantity_in(radii.unit)
return particles
def disable(self):
if self.is_supported():
self.conditions.code.disable_stopping_condition(self.type)
else:
name = [
name for name, value in self.conditions.all_conditions()
if value is self
][0]
raise AmuseException(
f"Can't disable stopping condition '{name}', since "
f"'{type(self.conditions.code).__name__}' does not "
"support this condition.")
def generate_variates(self, number_of_variates):
if self.base_distribution_type == "uniform":
return numpy.linspace(0.0,
1.0,
num=number_of_variates,
endpoint=False)
elif self.base_distribution_type == "random":
return numpy.random.uniform(0.0, 1.0, number_of_variates)
else:
raise AmuseException(
"Unknown base_distribution_type: {0}. Possible "
"options are: 'random' or 'uniform'.".format(
self.base_distribution_type))
def __init__(self,
grid,
number_of_sph_particles,
base_distribution_type="uniform",
seed=None):
if (grid.number_of_dimensions() != 3):
raise AmuseException("Grid must be 3D")
self.grid = grid
self.shape = grid.shape
self.number_of_sph_particles = number_of_sph_particles
self.base_distribution_type = base_distribution_type # "random" or "uniform"
if seed:
numpy.random.seed(seed)
def evolve_orbit_and_plot(codes, filename, parameters):
potential, cluster = setup_system()
orbits = []
for code in codes:
print("evolving with", code)
if code in codelist:
full_orbit = evolve_bridged_orbit_in_potential(
potential, cluster.copy(), code, **parameters)
else:
raise AmuseException("Unknow code: {}".format(code))
orbits.append(full_orbit)
print()
fit_orbit = load_fit_orbit()
plot_orbit(potential, orbits, codes, fit_orbit, filename)
def __init__(self,
grid,
number_of_sph_particles,
base_distribution_type="uniform",
seed=None):
if (grid.number_of_dimensions() != 3):
raise AmuseException("Grid must be 3D")
if not hasattr(grid, "momentum"):
grid.add_global_vector_attribute("momentum",
["rhovx", "rhovy", "rhovz"])
self.grid = grid
self.shape = grid.shape
self.number_of_sph_particles = number_of_sph_particles
self.base_distribution_type = base_distribution_type # "random" or "uniform"
if seed:
numpy.random.seed(seed)
def unpickle_stellar_structure(self):
if os.path.isfile(self.pickle_file):
infile = open(self.pickle_file, 'rb')
else:
raise AmuseException("Input pickle file '{0}' does not exist".format(self.pickle_file))
structure = pickle.load(infile)
self.mass = structure['mass']
self.radius = structure['radius']
self.number_of_zones = structure['number_of_zones']
self.number_of_species = structure['number_of_species']
self.species_names = structure['species_names']
self.density_profile = structure['density_profile']
self.radius_profile = structure['radius_profile']
self.mu_profile = structure['mu_profile']
self.composition_profile = structure['composition_profile']
self.specific_internal_energy_profile = structure['specific_internal_energy_profile']
self.midpoints_profile = structure['midpoints_profile']
def create_initial_wind_for_time(self, time, check_length=True):
"""
Particles are created as if the wind has already been blowing for
'time'. Note that this does not work if the mass loss is derived
from stellar evolution.
"""
self.model_time = time
self.particles.evolve_mass_loss(self.model_time)
if self.has_new_wind_particles():
wind_gas = self.create_wind_particles()
if self.has_target():
self.target_gas.add_particles(wind_gas)
elif check_length:
raise AmuseException("create_initial_wind time was too small to"
"create any particles.")
else:
wind_gas = Particles()
self.reset()
return wind_gas
def pynbody_column_density_plot(particles, width=None, qty='rho', units=None,
sideon=False, faceon=False, **kwargs):
if not HAS_PYNBODY:
raise AmuseException("Couldn't find pynbody")
if width is None:
width = 2.0 * particles.position.lengths_squared().amax().sqrt()
length_unit, pynbody_unit = _smart_length_units_for_pynbody_data(width)
pyndata = convert_particles_to_pynbody_data(particles, length_unit, pynbody_unit)
UnitlessArgs.strip([1]|length_unit, [1]|length_unit)
if sideon:
function = pynbody_sph.sideon_image
elif faceon:
function = pynbody_sph.faceon_image
else:
function = pynbody_sph.image
if units is None and qty == 'rho':
units = 'm_p cm^-2'
result = function(pyndata, width=width.value_in(length_unit), qty=qty, units=units, **kwargs)
UnitlessArgs.current_plot = native_plot.gca()
return result