def solar_system_in_time(time_JD=2457099.5|units.day): """ Initial conditions of Solar system -- particle set with the sun + eight planets, at the center-of-mass reference frame. Defined attributes: name, mass, radius, x, y, z, vx, vy, vz """ time_0 = 2457099.5 | units.day delta_JD = time_JD-time_0 sun, planets = get_sun_and_planets(delta_JD=delta_JD) solar_system = Particles() solar_system.add_particle(sun) solar_system.add_particles(planets) solar_system.move_to_center() ### to compare with JPL, relative positions and velocities need to be corrected for the # Sun's vectors with respect to the barycenter #r_s = (3.123390770608490E-03, -4.370830943817017E-04, -1.443425433116342E-04) | units.AU #v_s = (3.421633816761503E-06, 5.767414405893875E-06, -8.878039607570240E-08) | (units.AU / units.day) #print sun #print planets.position.in_(units.AU) + r_s #print planets.velocity.in_(units.AU/units.day) + v_s return solar_system
def sink_particles(sinks,sources,Raccretion=0.1 | units.AU):
closest=numpy.array([-1]*len(sources))
mind2=(numpy.array([Raccretion.number]*len(sources)) | Raccretion.unit)**2
for i,s in enumerate(sinks):
xs,ys,zs=s.x,s.y,s.z
d2=(sources.x-xs)**2+(sources.y-ys)**2+(sources.z-zs)**2
select=numpy.where( d2<mind2 )[0]
mind2[select]=d2[select]
closest[select]=i
to_remove=Particles(0)
for i,s in enumerate(sinks):
insink=numpy.where(closest == i)[0]
if len(insink) > 0:
cm=s.position*s.mass
p=s.velocity*s.mass
insinkp=Particles(0)
for ip in insink:
insinkp.add_particle(sources[ip])
s.mass+=insinkp.total_mass()
s.position=(cm+insinkp.center_of_mass()*insinkp.total_mass())/s.mass
s.velocity=(p+insinkp.total_momentum())/s.mass
# we lose angular momentum !
to_remove.add_particles(insinkp)
print len(insinkp),"particles accrete on star", i
if len(to_remove)>0:
sources.remove_particles(to_remove)
def get_stars_from_molecular_clous(parts):
cutoff_density = 10000 | units.amu/units.cm**3
stars = Particles(0)
for ip in parts:
if ip.rho>cutoff_density:
local_stars = make_stars(ip)
stars.add_particles(local_stars)
return stars
def get_stars_from_molecular_clous(parts):
cutoff_density = 10000 | units.amu / units.cm**3
stars = Particles(0)
for ip in parts:
if ip.rho > cutoff_density:
local_stars = make_stars(ip)
stars.add_particles(local_stars)
return stars
def get_potential_at_point(self,radius,x,y,z):
part=Particles(0)
for s in self.systems:
part.add_particles(s.particles)
phi=AdaptingVectorQuantity()
for rr,xx,yy,zz in zip(radius,x,y,z):
dr2=((part.x-xx)**2+(part.y-yy)**2+(part.z-zz)**2+rr**2+part.radius**2)
phi.append( (-self.G*part.mass/dr2**0.5).sum() )
return phi
def create_wind_particles(self):
wind = Particles(0)
for star in self.particles:
if star.lost_mass > self.sph_particle_mass:
new_particles = self.create_wind_particles_for_one_star(star)
wind.add_particles(new_particles)
star.wind_release_time = self.model_time
return wind
def accrete_gas(sink, gas):
"Accrete gas within sink radius"
accreted_gas = Particles()
distance_to_sink_squared = (
(gas.x - sink.x)**2
+ (gas.y - sink.y)**2
+ (gas.z - sink.z)**2
)
accreted_gas.add_particles(gas[
numpy.where(distance_to_sink_squared < sink.radius**2)
])
return accreted_gas
def get_gravity_at_point(self,radius,x,y,z):
part=Particles(0)
for s in self.systems:
part.add_particles(s.particles)
ax=AdaptingVectorQuantity()
ay=AdaptingVectorQuantity()
az=AdaptingVectorQuantity()
for rr,xx,yy,zz in zip(radius,x,y,z):
dr2=((part.x-xx)**2+(part.y-yy)**2+(part.z-zz)**2+rr**2+part.radius**2)
ax.append( (self.G*part.mass*(part.x-xx)/dr2**1.5).sum() )
ay.append( (self.G*part.mass*(part.y-yy)/dr2**1.5).sum() )
az.append( (self.G*part.mass*(part.z-zz)/dr2**1.5).sum() )
return ax,ay,az
def initialize_globular_clusters(cluster_population, nstars):
stars = Particles(0)
for ci in cluster_population:
converter = nbody_system.nbody_to_si(ci.mass, ci.radius)
bodies = new_king_model(nstars, ci.King_W0, converter)
bodies.parent = ci
bodies.position += ci.position
bodies.velocity += ci.velocity
bodies.name = "star"
ci.mass = bodies.mass.sum()
bodies.scale_to_standard(converter)
stars.add_particles(bodies)
return stars
def make_secondaries(center_of_masses, Nbin):
resulting_binaries = Particles()
singles_in_binaries = Particles()
binaries = center_of_masses.random_sample(Nbin)
mmin = center_of_masses.mass.min()
for bi in binaries:
mp = bi.mass
ms = numpy.random.uniform(mmin.value_in(units.MSun),
mp.value_in(units.MSun)) | units.MSun
a = random_semimajor_axis_PPE(mp, ms)
e = numpy.sqrt(numpy.random.random())
nb = new_binary_orbit(mp, ms, a, e)
nb.position += bi.position
nb.velocity += bi.velocity
nb = singles_in_binaries.add_particles(nb)
nb.radius = 0.01 * a
bi.radius = 3*a
binary_particle = bi.copy()
binary_particle.child1 = nb[0]
binary_particle.child2 = nb[1]
binary_particle.semi_major_axis = a
binary_particle.eccentricity = e
resulting_binaries.add_particle(binary_particle)
single_stars = center_of_masses-binaries
return single_stars, resulting_binaries, singles_in_binaries
def initialize_globular_clusters(cluster_population, nstars):
mmean = new_kroupa_mass_distribution(1000).sum() / 1000.
stars = Particles(0)
for ci in cluster_population:
nstars = max(1, int(ci.mass / mmean))
masses = new_kroupa_mass_distribution(nstars)
converter = nbody_system.nbody_to_si(masses.sum(), ci.radius)
bodies = new_king_model(nstars, ci.King_W0, converter)
bodies.parent = ci
bodies.mass = masses
bodies.position += ci.position
bodies.velocity += ci.velocity
bodies.name = "star"
ci.mass = bodies.mass.sum()
bodies.scale_to_standard(converter)
stars.add_particles(bodies)
return stars
def convert_stars(self, particles, stellar_evolution_code):
n_particles = self.divide_number_of_particles(particles)
se_colliders = particles.get_intersecting_subset_in(stellar_evolution_code.particles)
if self.verbose:
print("Converting stars of {0} to SPH models of {1} particles, respectively.".format(particles.mass, n_particles))
sph_models = (
self.relax(convert_stellar_model_to_SPH(se_colliders[0], n_particles[0], **self.star_to_sph_arguments)),
self.relax(convert_stellar_model_to_SPH(se_colliders[1], n_particles[1], **self.star_to_sph_arguments))
)
gas_particles = Particles()
for particle, sph_model in zip(particles, sph_models):
sph_model.position += particle.position
sph_model.velocity += particle.velocity
gas_particles.add_particles(sph_model)
if self.verbose:
print("Converting stars to SPH particles done")
if self.debug:
print(gas_particles)
return gas_particles
def all(self):
parts = Particles()
for parent in self:
if parent.subsystem is None:
parts.add_particle(parent)
else:
subsys = parts.add_particles(parent.subsystem)
subsys.position += parent.position
subsys.velocity += parent.velocity
return parts
def convert_stars(self, particles, stellar_evolution_code):
n_particles = self.divide_number_of_particles(particles)
se_colliders = particles.get_intersecting_subset_in(stellar_evolution_code.particles)
if self.verbose:
print "Converting stars of {0} to SPH models of {1} particles, respectively.".format(particles.mass, n_particles)
sph_models = (
self.relax(convert_stellar_model_to_SPH(se_colliders[0], n_particles[0], **self.star_to_sph_arguments)),
self.relax(convert_stellar_model_to_SPH(se_colliders[1], n_particles[1], **self.star_to_sph_arguments))
)
gas_particles = Particles()
for particle, sph_model in zip(particles, sph_models):
sph_model.position += particle.position
sph_model.velocity += particle.velocity
gas_particles.add_particles(sph_model)
if self.verbose:
print "Converting stars to SPH particles done"
if self.debug:
print gas_particles
return gas_particles
def plummer_w_binaries(N, f_bin=1.,logamin=-4,logamax=-0.5,seed=123454321):
random.seed(seed)
plummer=new_plummer_model(N)
semi=10**random.uniform(logamin,logamax,N) | nbody_system.length
inc=numpy.arccos(random.uniform(-1.,1.,N))/numpy.pi*180
longi=random.uniform(0,2*numpy.pi,N)/numpy.pi*180
binaries=Particles()
tobin=plummer[:int(f_bin*N)]
nobin=plummer-tobin
if len(nobin)>0:
binaries.add_particles(nobin)
for i,p in enumerate( tobin ):
mass1=p.mass/2
mass2=p.mass/2
binary=new_binary_from_orbital_elements(mass1,mass2,semi[i],
inclination=inc[i],longitude_of_the_ascending_node=longi[i])
binary.position+=p.position
binary.velocity+=p.velocity
binaries.add_particles(binary)
return binaries
def old_new_solar_system():
"""
Create initial conditions describing the solar system. Returns a single
particle set containing the sun, planets and Pluto. The model is centered at
the origin (center-of-mass(-velocity) coordinates).
Defined attributes:
name, mass, radius, x, y, z, vx, vy, vz
"""
sun = Particle()
sun.name = 'SUN'
sun.mass = 1.0 | units.MSun
sun.radius = 1.0 | units.RSun
planets = _planets_only()
particles = Particles()
particles.add_particle(sun)
particles.add_particles(planets)
particles.move_to_center()
return particles
def get_tt72_disk(m=10.e11|units.MSun,
r_min=25.|units.kpc,
n_rings=[12,15,18,21,24,27,30,33,36,39,42,45],
r_rings_rel=[0.2,0.25,0.3,0.35,0.4,0.45,0.5,0.55,0.6,0.65,0.7,0.75],
disk_id='a',
eps=0.|units.m):
"""
initialize disk ala TT72 (see the first paragraphs of sec. II and III of the paper)
positions and velocities with respect to the central particle of mass m
"""
disk = Particles()
for i,ri in enumerate(r_rings_rel):
disk_rad_i = Particles(n_rings[i])
a = ri*r_min
phi_i = numpy.linspace(0., pipi, num=n_rings[i], endpoint=False)
disk_rad_i.x = a * numpy.cos(phi_i)
disk_rad_i.y = a * numpy.sin(phi_i)
disk_rad_i.z = 0. * a
x_r = disk_rad_i.x/a
y_r = disk_rad_i.y/a
#vc = (constants.G*m/a)**0.5
vc = ( constants.G*m*a**2/(a**2 + eps**2)**1.5 )**0.5
disk_rad_i.vx = -vc * y_r
disk_rad_i.vy = vc * x_r
disk_rad_i.vz = 0.0 * vc
disk.add_particles(disk_rad_i)
# test particles
disk.mass = 0.|units.MSun
# identification of the disk
disk.id = disk_id
return disk
def run_instances_in_serial(self, end_time):
# final particle set of stones
stones_final = Particles(0)
# list of final binary particles sets
# should be the same binary number_of_subsets-times
list_binary_final = []
for i, instance_i in enumerate(self.list_of_instances):
instance_i.evolve_model(end_time)
stones_final.add_particles(instance_i.particles[2:].copy())
list_binary_final.append(instance_i.particles[:2].copy())
# check if binary evolved the same in all subsets
diff_flag = self.final_binary_diff(list_binary_final)
if diff_flag != 0:
print "\t binary of the first subset used for the evolution"
# final particle set after evolve
binary_and_stones_final = Particles(0)
binary_and_stones_final.add_particles(list_binary_final[0])
binary_and_stones_final.add_particles(stones_final)
return binary_and_stones_final
class StubInterface(object):
def __init__(self, **options):
self.maximum_density = 1 | units.kg / units.m**3
self._gas_particles = Particles()
self._dm_particles = Particles()
self._all_particles = ParticlesSuperset([self._gas_particles, self._dm_particles])
def before_get_parameter(self):
pass
def before_set_parameter(self):
pass
def initialize_code(self):
return 0
synchronize_model = commit_particles = recommit_particles = commit_parameters = initialize_code
def new_particle(self, mass, x, y, z, vx, vy, vz, *args):
next_id = len(self._dm_particles)
temp = Particles(len(mass))
temp.mass = mass
temp.x = x
temp.y = y
temp.z = z
temp.vx = vx
temp.vy = vy
temp.vz = vz
temp.id = list(range(next_id, next_id + len(mass)))
self._dm_particles.add_particles(temp)
return [temp.id, temp.id]
def new_gas_particle(self, mass, x, y, z, vx, vy, vz, *args):
next_id = len(self._gas_particles) + 1000000
temp = Particles(len(mass))
temp.mass = mass
temp.x = x
temp.y = y
temp.z = z
temp.vx = vx
temp.vy = vy
temp.vz = vz
temp.id = list(range(next_id, next_id + len(mass)))
self._gas_particles.add_particles(temp)
return [temp.id, temp.id]
def delete_particle(self, indices):
for index in indices:
for id, particle in zip(self._all_particles.id, self._all_particles):
if id == index:
self._all_particles.remove_particle(particle)
return 0
def get_mass(self, indices):
return [[mass for index in indices for id, mass in zip(self._all_particles.id,
self._all_particles.mass) if index == id], [0]*len(indices)]
def set_mass(self, indices, masses):
for index, mass in zip(indices, masses):
for id, particle in zip(self._all_particles.id, self._all_particles):
if id == index:
particle.mass = mass
break
return 0
def get_position(self, indices):
return [[x for index in indices for id, x in zip(self._all_particles.id, self._all_particles.x) if index == id],
[y for index in indices for id, y in zip(self._all_particles.id, self._all_particles.y) if index == id],
[z for index in indices for id, z in zip(self._all_particles.id, self._all_particles.z) if index == id],
[0]*len(indices)]
def get_velocity(self, indices):
return [[vx for index in indices for id, vx in zip(self._all_particles.id, self._all_particles.vx) if index == id],
[vy for index in indices for id, vy in zip(self._all_particles.id, self._all_particles.vy) if index == id],
[vz for index in indices for id, vz in zip(self._all_particles.id, self._all_particles.vz) if index == id],
[0]*len(indices)]
def has_stopping_condition(self, type):
return 1 if type == 6 else 0
def get_stopping_condition_maximum_density_parameter(self):
return self.maximum_density
def set_stopping_condition_maximum_density_parameter(self, value):
self.maximum_density = value
is_stopping_condition_set = is_stopping_condition_enabled = has_stopping_condition
def get_number_of_stopping_conditions_set(self):
return 3
def get_stopping_condition_info(self, sc_indices):
return [6]*len(sc_indices), [1]*len(sc_indices)
def get_stopping_condition_particle_index(self, sc_index, sc_sub_index):
return range(len(self._gas_particles) + 1000000 - len(sc_index), len(self._gas_particles) + 1000000)
def enable_stopping_condition(self, type):
pass
def evolve_model(self, time):
return 0
def handle_collisions(self, primaries, secondaries):
result = Particles()
for primary, secondary in zip(primaries.as_set(),
secondaries.as_set()):
result.add_particles(self.handle_collision(primary, secondary))
return result
class StubInterface(object):
def __init__(self, **options):
self.maximum_density = 1 | units.kg / units.m**3
self._gas_particles = Particles()
self._dm_particles = Particles()
self._all_particles = ParticlesSuperset([self._gas_particles, self._dm_particles])
def before_get_parameter(self):
pass
def before_set_parameter(self):
pass
def initialize_code(self):
return 0
synchronize_model = commit_particles = recommit_particles = commit_parameters = initialize_code
def new_particle(self, mass, x, y, z, vx, vy, vz, *args):
next_id = len(self._dm_particles)
temp = Particles(len(mass))
temp.mass = mass
temp.x = x
temp.y = y
temp.z = z
temp.vx = vx
temp.vy = vy
temp.vz = vz
temp.id = list(range(next_id, next_id + len(mass)))
self._dm_particles.add_particles(temp)
return [temp.id, temp.id]
def new_gas_particle(self, mass, x, y, z, vx, vy, vz, *args):
next_id = len(self._gas_particles) + 1000000
temp = Particles(len(mass))
temp.mass = mass
temp.x = x
temp.y = y
temp.z = z
temp.vx = vx
temp.vy = vy
temp.vz = vz
temp.id = list(range(next_id, next_id + len(mass)))
self._gas_particles.add_particles(temp)
return [temp.id, temp.id]
def delete_particle(self, indices):
for index in indices:
for id, particle in zip(self._all_particles.id, self._all_particles):
if id == index:
self._all_particles.remove_particle(particle)
return 0
def get_mass(self, indices):
return [[mass for index in indices for id, mass in zip(self._all_particles.id,
self._all_particles.mass) if index == id], [0]*len(indices)]
def set_mass(self, indices, masses):
for index, mass in zip(indices, masses):
for id, particle in zip(self._all_particles.id, self._all_particles):
if id == index:
particle.mass = mass
break
return 0
def get_position(self, indices):
return [[x for index in indices for id, x in zip(self._all_particles.id, self._all_particles.x) if index == id],
[y for index in indices for id, y in zip(self._all_particles.id, self._all_particles.y) if index == id],
[z for index in indices for id, z in zip(self._all_particles.id, self._all_particles.z) if index == id],
[0]*len(indices)]
def get_velocity(self, indices):
return [[vx for index in indices for id, vx in zip(self._all_particles.id, self._all_particles.vx) if index == id],
[vy for index in indices for id, vy in zip(self._all_particles.id, self._all_particles.vy) if index == id],
[vz for index in indices for id, vz in zip(self._all_particles.id, self._all_particles.vz) if index == id],
[0]*len(indices)]
def has_stopping_condition(self, type):
return 1 if type == 6 else 0
def get_stopping_condition_maximum_density_parameter(self):
return self.maximum_density
def set_stopping_condition_maximum_density_parameter(self, value):
self.maximum_density = value
is_stopping_condition_set = is_stopping_condition_enabled = has_stopping_condition
def get_number_of_stopping_conditions_set(self):
return 3
def get_stopping_condition_info(self, sc_indices):
return [6]*len(sc_indices), [1]*len(sc_indices)
def get_stopping_condition_particle_index(self, sc_index, sc_sub_index):
return list(range(len(self._gas_particles) + 1000000 - len(sc_index), len(self._gas_particles) + 1000000))
def enable_stopping_condition(self, type):
pass
def evolve_model(self, time):
return 0
def run_diagnostics(
model,
logger=None,
length_unit=units.pc,
mass_unit=units.MSun,
time_unit=units.Myr,
):
"""
Run diagnostics on model
"""
logger = logger or logging.getLogger(__name__)
stars = model.star_particles
sinks = model.sink_particles
gas = model.gas_particles
converter = model.star_converter
if not sinks.is_empty():
non_collisional_bodies = Particles()
non_collisional_bodies.add_particles(stars)
non_collisional_bodies.add_particles(sinks)
else:
non_collisional_bodies = stars
groups = identify_subgroups(
converter,
non_collisional_bodies,
peak_density_threshold=1e-16 | units.g * units.cm**-3,
)
n_groups = len(groups)
if hasattr(stars, 'group_id'):
group_id_offset = 1 + max(stars.group_id)
else:
group_id_offset = 1 # a group id of 0 would mean "no group found"
logger.info("Found %i groups", n_groups)
for i, group in enumerate(groups):
group_id = i + group_id_offset
if hasattr(group, 'group_id'):
group.previous_group_id = group.group_id
else:
group.previous_group_id = 0
group.group_id = group_id
stars_in_group = len(group)
if (stars_in_group > 100):
mass_in_group = group.total_mass().in_(mass_unit)
mass_fraction = [0.01, 0.02, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]
radii, new_mass_fraction = group.LagrangianRadii(
unit_converter=converter,
mf=mass_fraction,
cm=group.center_of_mass(),
)
assert (new_mass_fraction == mass_fraction)
radii = radii.value_in(length_unit)
x, y, z = group.center_of_mass().value_in(length_unit)
median_previous_group_id = numpy.median(group.previous_group_id)
logger.info(
"step %i group %i nstars %i mass %s xyz %f %f %f %s origin %i "
"LR %f %f %f %f %f %f %f %f %f %s",
model.step,
group_id,
stars_in_group,
mass_in_group,
x,
y,
z,
length_unit,
median_previous_group_id,
radii[0],
radii[1],
radii[2],
radii[3],
radii[4],
radii[5],
radii[6],
radii[7],
radii[8],
length_unit,
)
groups = ParticlesSuperset(groups)
group_identifiers = Particles(keys=groups.key)
group_identifiers.group_id = groups.group_id
group_identifiers.previous_group_id = groups.previous_group_id
return group_identifiers
def new_binary_distribution(
primary_mass,
secondary_mass=None,
binaries=None,
min_mass=0.08 | units.MSun,
):
"""
Takes primary masses, and returns a set of stars and a set of binaries
formed by these stars.
Secondary masses are given by a uniform random mass ratio with the primary
masses. If optional secondary masses are given, these are used instead.
binaries is an optional particleset used for the binary pairs, with given
positions and velocities. Other parameters are ignored.
"""
N = len(primary_mass)
if binaries is None:
binaries = Particles(N)
# Should give some position/velocity as well?
elif len(binaries) != N:
print("binaries must be None or have the same lenght as primary_mass")
return -1
if secondary_mass is None:
# Now, we need to specify the mass ratio in the binaries.
# A flat distribution seems to be OK.
mass_ratio = uniform(N)
# This gives us the secondaries' masses
secondary_mass = mass_ratio * primary_mass
# secondaries are min_mass at least
secondary_mass = numpy.maximum(secondary_mass, min_mass)
elif len(secondary_mass) != N:
print("Number of secondaries is unequal to number of primaries!")
return -1
else:
# Make sure primary_mass is the larger of the two, and secondary_mass
# the smaller.
pm = primary_mass.maximum(secondary_mass)
sm = primary_mass.minimum(secondary_mass)
primary_mass = pm
secondary_mass = sm
del(pm, sm)
# Now, we need to calculate the semi-major axes for the binaries. Since the
# observed quantity is orbital periods, we start from there.
mean_log_orbital_period = 5 # 10log of the period in days, (Duchene&Kraus)
sigma_log_orbital_period = 2.3
orbital_period = numpy.random.lognormal(
size=N,
mean=numpy.log(10) * mean_log_orbital_period,
sigma=numpy.log(10) * sigma_log_orbital_period,
) | units.day
# We need the masses to calculate the corresponding semi-major axes.
semi_major_axis = orbital_period_to_semi_major_axis(
orbital_period,
primary_mass,
secondary_mass,
)
# Eccentricity: square root of random value
eccentricity = numpy.sqrt(random(N))
# Other orbital elements at random
inclination = pi * random(N) | units.rad
true_anomaly = 2 * pi * random(N) | units.rad
longitude_of_the_ascending_node = 2 * pi * random(N) | units.rad
argument_of_periapsis = 2 * pi * random(N) | units.rad
primaries, secondaries = generate_binaries(
primary_mass,
secondary_mass,
semi_major_axis,
eccentricity=eccentricity,
inclination=inclination,
true_anomaly=true_anomaly,
longitude_of_the_ascending_node=longitude_of_the_ascending_node,
argument_of_periapsis=argument_of_periapsis,
G=constants.G,
)
stars = Particles()
primaries.position += binaries.position
secondaries.position += binaries.position
primaries.velocity += binaries.velocity
secondaries.velocity += binaries.velocity
primaries = stars.add_particles(primaries)
secondaries = stars.add_particles(secondaries)
binaries.eccentricity = eccentricity
binaries.semi_major_axis = semi_major_axis
for i in range(len(primaries)):
binaries[i].child1 = primaries[i]
binaries[i].child2 = secondaries[i]
# Probably needed
binaries[i].mass = primaries[i].mass + secondaries[i].mass
return stars, binaries
def iliev_test_7_ic(N=10000, Ns=10, L=6.6 | units.kpc):
print "Initializing iliev_test_7"
mp = rhoinit * (2 * L)**3 / N
Nc = ((rhoclump * 4 * constants.pi / 3 * (0.8 | units.kpc)**3) / mp)
Nc = int(Nc)
print Nc
try:
f = open("glass%9.9i.pkl" % N, "rb")
x, y, z = cPickle.load(f)
f.close()
except:
x, y, z = glass_unit_cube(N, target_rms=0.05).make_xyz()
f = open("glass%9.9i.pkl" % N, "wb")
cPickle.dump((x, y, z), f)
f.close()
sel = numpy.where(((x -
(5. / 6.6))**2 + y**2 + z**2)**0.5 > (0.8 / 6.6))[0]
x = x[sel]
y = y[sel]
z = z[sel]
p = Particles(len(x))
print len(x)
# set particles homogeneously in space
p.x = L * x
p.y = L * y
p.z = L * z
# set other properties
p.h_smooth = 0. | units.parsec
p.vx = 0. | (units.km / units.s)
p.vy = 0. | (units.km / units.s)
p.vz = 0. | (units.km / units.s)
p.u = uinit
p.rho = rhoinit
p.mass = mp
p.flux = 0. | (units.s**-1)
p.xion = 0. | units.none
sources = Particles(Ns)
x, y, z = uniform_unit_sphere(Ns).make_xyz()
if Ns == 1:
x, y, z = 0., 0., 0.
sources.x = -(5. | units.kpc) + L * x * (1. / N)**(1. / 3) / 10
sources.y = L * y * (1. / N)**(1. / 3) / 10
sources.z = L * z * (1. / N)**(1. / 3) / 10
sources.luminosity = (1.2e52 / Ns) | (units.s**-1)
sources.SpcType = 1.
clump = Particles(Nc)
x, y, z = uniform_unit_sphere(Nc).make_xyz()
clump.x = (5. | units.kpc) + (0.8 | units.kpc) * x
clump.y = (0.8 | units.kpc) * y
clump.z = (0.8 | units.kpc) * z
clump.h_smooth = 0. | units.parsec
clump.vx = 0. | (units.km / units.s)
clump.vy = 0. | (units.km / units.s)
clump.vz = 0. | (units.km / units.s)
clump.u = uclump
clump.rho = rhoclump
clump.mass = mp
clump.flux = 0. | (units.s**-1)
clump.xion = 0. | units.none
p.add_particles(clump)
return p, sources
def run(self, run_num):
print("Setting up run %d..." % (run_num))
# Create group for run
run_group = self.hdf5_file.create_group(str(run_num))
# create nbody converter thing?
convert_nbody = nbody_system.nbody_to_si(100.0 | units.MSun, 1 | units.parsec)
# initialize particle datamodel
stars = Particles()
print(" Initializing and populating clusters...")
# populate/initialize clusters, add to particle model
for c in self.clusters:
c.populate()
stars.add_particles(c.plummer)
print(" Done!")
# initialize codes
# initialize hermite
print(" Initializing hermite...")
hermite_code = Hermite(convert_nbody)
hermite_code.particles.add_particles(stars)
detect_coll = hermite_code.stopping_conditions.collision_detection
detect_coll.enable()
print(" Done!")
print(" Initializing SSE...")
sse_code = SSE()
sse_code.particles.add_particles(stars)
print("Done!\n")
# actually run the simulation!
print("===== Run #%d =====" % (run_num))
t = 0 # time
i = 1 # iteration num
c = 1 # collision num
while(t < self.runtime):
print(" Time (Myr): %d" % (t))
print(" Simulating...")
# evolve model
hermite_code.evolve_model(t | units.Myr)
sse_code.evolve_model(t | units.Myr)
hermite_code.particles.copy_values_of_attribute_to("position", sse_code.particles)
sse_code.particles.synchronize_to(hermite_code.particles)
if detect_coll.is_set():
print("Detected a collision!!")
print(detect_coll.particles(0))
coll_f = open(sub_folder + "/collision-" + str(c) + "_time" + str(t) + ".txt")
coll_f.write( detect_coll.particles(0) )
coll_f.close()
c += 1
# somehow put a log
# also... eventually, "pause" the rest of the simulation and simulate the collision somehow
#sse_code.particles.copy_values_of_attribute_to("mass", hermite_code.particles)
#sse_code.particles.copy_values_of_attribute_to("stellar_type", hermite_code.particles)
#sse_code.particles.copy_values_of_attribute_to("age", hermite_code.particles)
#sse_code.particles.copy_values_of_attribute_to("luminosity", hermite_code.particles)
#sse_code.particles.copy_values_of_attribute_to("temperature", hermite_code.particles)
#sse_code.particles.copy_values_of_attribute_to("radius", hermite_code.particles)
hermite_code.particles.mass = sse_code.particles.mass
print(" Done.")
# output to csv
print(" Outputting to HDF5...")
#file_name = sub_folder + "/data-" + str(i) + "_time-" + str(t) + ".txt"
self.write_to_hdf5_file(run_group, i, t, stars)
print(" Done.")
#print(" Creating plot...")
#plot_name = "System at " + str(t) + " Myr"
#plot_path = sub_folder + "/plot-" + str(i) + "_time-" + str(t) + ".png"
#plot_data(plot_name, plot_path, hermite_code.particles)
#print(" Done.")
t += self.timestep
i += 1
print("Run complete.\n")
def handle_collisions(self, primaries, secondaries):
result = Particles()
for primary, secondary in zip(primaries.as_set(), secondaries.as_set()):
result.add_particles(self.handle_collision(primary, secondary))
return result
class MinimalWorkingExample(object):
def __init__(self,
total_N=16000, #total number of disc particles
tend=250. | units.yr, #End time of the simulation
Mstar=1. | units.MSun, #Mass of the accreting star
dt = 1. | units.yr, #timestep
temp = 25. | units.K, #temperature of the ISM
v_ism = 3.0 | units.kms, #Velocity of the ISM with respect to the star
n_core = 4, #Number of cores used by the community codes
mu = 2.3, #mean molecular weight
theta = 0., #angle of the disc with respect to the flow, 0 means disc rotation axis is parallel with flow direction
n_dens = 5e6 | (units.cm)**(-3.), #Number density of the ism
dirname = './'):
self.disc_N = total_N
self.tend = tend
self.Mstar = Mstar
self.dt = dt
self.mu = mu * constants.proton_mass
self.T = temp
self.cs = ((temp*constants.kB)/self.mu).sqrt()
# print "cs=", self.cs.in_(units.kms)
self.v_ism = v_ism
self.ism_n_dens = n_dens #number density of the ism
self.ism_dens = self.mu*self.ism_n_dens #mass density of the ism
self.converter = nbody_system.nbody_to_si(self.Mstar, 1. | units.AU)
self.disc_angle = numpy.radians(theta)
self.n_core = n_core
self.dirname = dirname
self.discfraction = 0.01
self.disc_Rmin = self.converter.to_generic(10.| units.AU).value_in(nbody_system.length)
self.disc_Rmax = self.converter.to_generic(100.| units.AU).value_in(nbody_system.length)
self.filename = "DiskWind.h5"
print('Files will be saved in ', self.dirname)
def initialize_star(self):
self.star=Particles(1)
self.star.mass=self.Mstar
self.star.radius= 1. | units.AU
self.star.x=0.|units.AU
self.star.y=0.|units.AU
self.star.z=0.|units.AU
self.star.vx=0.|units.kms
self.star.vy=0.|units.kms
self.star.vz=0.|units.kms
def initialize_data(self):
self.cylinder_radius = 500. | units.AU
self.cylinder_length = 2.*self.cylinder_radius
self.cylinder_vol = self.cylinder_length*numpy.pi*self.cylinder_radius**2.
self.sph_particle_dens = self.ism_dens/(0.01*self.Mstar/self.disc_N)
self.initialize_star()
self.initialize_ism()
def initialize_ism(self):
#specific internal energy of the ism, i.e. internal energy per unit mass
#In case of an isothermal EOS, Gadget requires the sound-speed squared as input parameter instead of the thermal energy per unit mass
self.ism_u = self.cs**2.
self.ism_slice_length = self.v_ism*self.dt
self.ism_slice_vol = self.ism_slice_length*numpy.pi*self.cylinder_radius**2.0
self.ism_slice_mass = self.ism_dens*self.ism_slice_vol
self.ism_mass = self.ism_dens*self.cylinder_vol
#The x-coordinate of the inflow
self.offset = self.cylinder_radius
self.ism_slice_N = int(round(self.sph_particle_dens*self.ism_slice_vol))
self.total_N = self.ism_slice_N*(self.cylinder_length/self.ism_slice_length)
def make_slice(self):
new_slice = Particles(self.ism_slice_N)
rho = numpy.sqrt(numpy.random.uniform(0,1,self.ism_slice_N))*self.cylinder_radius.value_in(units.AU)
phi = numpy.random.uniform(0,2.0*numpy.pi, self.ism_slice_N)
new_slice.x = numpy.random.uniform(self.ism_slice_length.value_in(units.AU), 0, self.ism_slice_N) - self.offset.value_in(units.AU) | units.AU
new_slice.y = rho*numpy.sin(phi) | units.AU
new_slice.z = rho*numpy.cos(phi) | units.AU
new_slice.vx = 1000*self.v_ism
new_slice.vy = 0.0 | units.kms
new_slice.vz = 0.0 | units.kms
new_slice.mass = self.ism_slice_mass/self.ism_slice_N
new_slice.u = self.ism_u
return new_slice
def create_disc(self):
#The following line makes sure masses for ISM and disc particles are equal:
self.discfraction = (self.disc_N*(self.ism_slice_mass/self.ism_slice_N))/self.Mstar
T_disc = self.T
cs_disc = ((T_disc*constants.kB)/self.mu).sqrt()
densitypower = 1.5
g2=2-densitypower
k_out=((1+self.discfraction)/self.disc_Rmax**3)**0.5
sigma_out=g2*self.discfraction/(2*numpy.pi*self.disc_Rmax**densitypower*(self.disc_Rmax**g2-self.disc_Rmin**g2))
q_out = self.converter.to_generic(cs_disc).value_in(nbody_system.length/nbody_system.time)/(numpy.pi*sigma_out/k_out)
print("Number of disk particles:", self.disc_N)
proto=ProtoPlanetaryDisk(self.disc_N,
convert_nbody=self.converter,
discfraction=self.discfraction,
densitypower=1.5,
thermalpower=0,
Rmin=self.disc_Rmin,
Rmax=self.disc_Rmax,
q_out=q_out)
disc=proto.result
print("The mass of a disc particle = ", disc.mass[0].value_in(units.kg))
#Rotate 90 degrees with respect to the z-axis and then theta degrees with respect to the y-axis
temp_x = disc[:].x
temp_y = disc[:].y
temp_z = disc[:].z
temp_vx = disc[:].vx
temp_vy = disc[:].vy
temp_vz = disc[:].vz
disc.x = temp_z*numpy.cos(self.disc_angle) - temp_y*numpy.sin(self.disc_angle)
disc.y = temp_z*numpy.sin(self.disc_angle) + temp_y*numpy.cos(self.disc_angle)
disc.z = -temp_x
disc.vx = temp_vz*numpy.cos(self.disc_angle) - temp_vy*numpy.sin(self.disc_angle)
disc.vy = temp_vz*numpy.sin(self.disc_angle) + temp_vy*numpy.cos(self.disc_angle)
disc.vz = -temp_vx
return disc
def evolve_model(self):
self.initialize_data()
self.ism_code = Gadget2(self.converter, number_of_workers=self.n_core)#, debugger='gdb')
self.ism_code.parameters.time_max = 1024*self.dt
self.ism_code.parameters.n_smooth = 64
self.ism_code.parameters.n_smooth_tol = 2./64.
self.ism_code.parameters.artificial_viscosity_alpha = 0.1
self.ism_code.parameters.epsilon_squared = (1. | units.AU)**2.
self.all_particles = Particles()
write_set_to_file(self.star, self.filename, "hdf5", append_to_file=False)
write_set_to_file(self.all_particles, self.filename, "hdf5", append_to_file=False)
self.initial_disc_particles = self.create_disc()
self.all_particles.add_particles(self.initial_disc_particles)
self.ism_code.gas_particles.add_particles(self.initial_disc_particles)
#You can only add a sink after adding gas particles
#starinsph refers to the corresponding particle set/id in the community code
starinsph = self.ism_code.dm_particles.add_particles(self.star)
#Use the build-in sink particle routine from amuse.ext.sink. Sink_radius needs to be defined manually otherwise the particle radius in gadget is taken,
#which does not corresponding to the particle radius in the framework (since 'particles' in gadget do not have a radius, it is set to 0.01 | generic_unit_system.length
#and corresponds to the gas gravitational smoothing epsilon.
sink = new_sink_particles(starinsph, sink_radius= self.star.radius)
self.channel_from_ismcode_to_framework = self.ism_code.gas_particles.new_channel_to(self.all_particles)
time = 0. | units.yr
while time <= (self.tend+self.dt/2.):
print("Adding new slice of ISM...")
newslice=self.make_slice()
self.ism_code.gas_particles.add_particles(newslice)
self.all_particles.add_particles(newslice)
start = timing.time()
print("=======================================================")
print("Evolving to time = ", time.value_in(units.yr), " of ", self.tend.value_in(units.yr)," years...")
self.ism_code.evolve_model(time)
print("This took ", (timing.time() - start), " s")
out_of_bounds = self.ism_code.gas_particles.select_array(lambda x,y,z:(x > self.cylinder_radius)|((z**2+y**2).sqrt() >= self.cylinder_radius), ["x","y","z"])
if len(out_of_bounds)>0:
print("Removing ", len(out_of_bounds), " particles from the code because they were out of bounds")
self.ism_code.gas_particles.remove_particles(out_of_bounds)
self.ism_code.gas_particles.synchronize_to(self.all_particles)
sink.accrete(self.ism_code.gas_particles)
write_set_to_file(self.star, self.filename, "hdf5")
write_set_to_file(self.all_particles, self.filename, "hdf5")
time += self.dt
print("=======================================================")
self.ism_code.stop()
def form_stars_from_group_older_version(
group_index,
sink_particles,
newly_removed_gas,
lower_mass_limit=settings.stars_lower_mass_limit,
upper_mass_limit=settings.stars_upper_mass_limit,
local_sound_speed=0.2 | units.kms,
minimum_sink_mass=0.01 | units.MSun,
logger=None,
randomseed=None,
shrink_sinks=True,
**keyword_arguments):
"""
Form stars from specific group of sinks.
NOTE: This is the older version where removed gas is
considered as star-forming region. This is now being
updated to the above latest version.
"""
logger = logger or logging.getLogger(__name__)
logger.info("Using form_stars_from_group on group %i", group_index)
if randomseed is not None:
logger.info("Setting random seed to %i", randomseed)
numpy.random.seed(randomseed)
# Sanity check: each sink particle must be in a group.
ungrouped_sinks = sink_particles.select_array(lambda x: x <= 0,
['in_group'])
if not ungrouped_sinks.is_empty():
logger.info(
"WARNING: There exist ungrouped sinks. Something is wrong!")
return None
# Consider only group with input group index from here onwards.
group = sink_particles[sink_particles.in_group == group_index]
# Sanity check: group must have at least a sink
if group.is_empty():
logger.info(
"WARNING: There is no sink in the group: Something is wrong!")
return None
number_of_sinks = len(group)
logger.info("%i sinks found in group #%i: %s", number_of_sinks,
group_index, group.key)
group_mass = group.total_mass()
logger.info("Group mass: %s", group_mass.in_(units.MSun))
next_mass = generate_next_mass(
initial_mass_function=initial_mass_function,
lower_mass_limit=lower_mass_limit,
upper_mass_limit=upper_mass_limit,
)[0][0]
try:
# Within a group, group_next_primary_mass values are either
# a mass, or 0 MSun. If all values are 0 MSun, this is a
# new group. Else, only interested on the non-zero value. The
# non-zero values are the same.
logger.info('SANITY CHECK: group_next_primary_mass %s',
group.group_next_primary_mass)
if group.group_next_primary_mass.max() == 0 | units.MSun:
logger.info('Initiate group #%i for star formation', group_index)
group.group_next_primary_mass = next_mass
else:
next_mass = group.group_next_primary_mass.max()
# This happens for the first ever assignment of this attribute
except AttributeError:
logger.info(
'AttributeError exception: Initiate group #%i for star formation',
group_index)
group.group_next_primary_mass = next_mass
logger.info("Next mass is %s", next_mass)
if group_mass < next_mass:
logger.info("Group #%i is not massive enough for the next star",
group_index)
return None
# Form stars from the leftover group sink mass
mass_left = group_mass - next_mass
masses = new_masses(
stellar_mass=mass_left,
lower_mass_limit=lower_mass_limit,
upper_mass_limit=upper_mass_limit,
initial_mass_function=settings.stars_initial_mass_function,
)
number_of_stars = len(masses)
logger.info("%i stars created in group #%i with %i sinks", number_of_stars,
group_index, number_of_sinks)
new_stars = Particles(number_of_stars)
new_stars.age = 0 | units.Myr
new_stars[0].mass = next_mass
new_stars[1:].mass = masses[:-1]
group.group_next_primary_mass = masses[-1]
new_stars = new_stars.sorted_by_attribute("mass").reversed()
logger.info("Group's next primary mass is %s",
group.group_next_primary_mass[0])
# Create placeholders for attributes of new_stars
new_stars.position = [0, 0, 0] | units.pc
new_stars.velocity = [0, 0, 0] | units.kms
new_stars.origin_cloud = group[0].key
new_stars.star_forming_radius = 0 | units.pc
new_stars.star_forming_u = local_sound_speed**2
# Find the newly removed gas in the group
removed_gas = Particles()
if not newly_removed_gas.is_empty():
for s in group:
removed_gas_by_this_sink = (
newly_removed_gas[newly_removed_gas.accreted_by_sink == s.key])
removed_gas.add_particles(removed_gas_by_this_sink)
logger.info("%i removed gas found in this group", len(removed_gas))
# Star forming regions that contain the removed gas and the group
# of sinks
if not removed_gas.is_empty():
removed_gas.radius = removed_gas.h_smooth
star_forming_regions = group.copy()
star_forming_regions.density = (
star_forming_regions.initial_density / 1000
) # /1000 to reduce likelihood of forming stars in sinks
star_forming_regions.accreted_by_sink = star_forming_regions.key
try:
star_forming_regions.u = star_forming_regions.u
except AttributeError:
star_forming_regions.u = local_sound_speed**2
star_forming_regions.add_particles(removed_gas.copy())
star_forming_regions.sorted_by_attribute("density").reversed()
# Generate a probability list of star forming region indices the
# stars should associate to
probabilities = (star_forming_regions.density /
star_forming_regions.density.sum())
probabilities /= probabilities.sum() # Ensure sum is exactly 1
logger.info("Max & min probabilities: %s, %s", probabilities.max(),
probabilities.min())
logger.info("%i star forming regions", len(star_forming_regions))
def delta_positions_and_velocities(new_stars, star_forming_regions,
probabilities):
"""
Assign positions and velocities of stars in the star forming regions
according to the probability distribution
"""
number_of_stars = len(new_stars)
# Create an index list of removed gas from probability list
sample = numpy.random.choice(len(star_forming_regions),
number_of_stars,
p=probabilities)
# Assign the stars to the removed gas according to the sample
star_forming_regions_sampled = star_forming_regions[sample]
new_stars.position = star_forming_regions_sampled.position
new_stars.velocity = star_forming_regions_sampled.velocity
new_stars.origin_cloud = star_forming_regions_sampled.accreted_by_sink
new_stars.star_forming_radius = star_forming_regions_sampled.radius
try:
new_stars.star_forming_u = star_forming_regions_sampled.u
except AttributeError:
new_stars.star_forming_u = local_sound_speed**2
# Random position of stars within the sink radius they assigned to
rho = (numpy.random.random(number_of_stars) *
new_stars.star_forming_radius)
theta = (numpy.random.random(number_of_stars) *
(2 * numpy.pi | units.rad))
phi = (numpy.random.random(number_of_stars) * numpy.pi | units.rad)
x = (rho * sin(phi) * cos(theta)).value_in(units.pc)
y = (rho * sin(phi) * sin(theta)).value_in(units.pc)
z = (rho * cos(phi)).value_in(units.pc)
X = list(zip(*[x, y, z])) | units.pc
# Random velocity, sample magnitude from gaussian with local sound
# speed like Wall et al (2019)
# temperature = 10 | units.K
# or (gamma * local_pressure / density).sqrt()
velocity_magnitude = numpy.random.normal(
# loc=0.0, # <- since we already added the velocity of the sink
scale=new_stars.star_forming_u.sqrt().value_in(units.kms),
size=number_of_stars,
) | units.kms
velocity_theta = (numpy.random.random(number_of_stars) *
(2 * numpy.pi | units.rad))
velocity_phi = (numpy.random.random(number_of_stars) *
(numpy.pi | units.rad))
vx = (velocity_magnitude * sin(velocity_phi) *
cos(velocity_theta)).value_in(units.kms)
vy = (velocity_magnitude * sin(velocity_phi) *
sin(velocity_theta)).value_in(units.kms)
vz = (velocity_magnitude * cos(velocity_phi)).value_in(units.kms)
V = list(zip(*[vx, vy, vz])) | units.kms
return X, V
dX, dV = delta_positions_and_velocities(new_stars, star_forming_regions,
probabilities)
logger.info("Updating new stars...")
new_stars.position += dX
new_stars.velocity += dV
# For Pentacle, this is the PP radius
new_stars.radius = 0.05 | units.parsec
# mass_ratio = 1 - new_stars.total_mass()/group.total_mass()
# group.mass *= mass_ratio
excess_star_mass = 0 | units.MSun
for s in group:
logger.info('Sink mass before reduction: %s', s.mass.in_(units.MSun))
total_star_mass_nearby = (
new_stars[new_stars.origin_cloud == s.key]).total_mass()
# To prevent sink mass becomes negative
if s.mass > minimum_sink_mass:
if (s.mass - total_star_mass_nearby) <= minimum_sink_mass:
excess_star_mass += (total_star_mass_nearby - s.mass +
minimum_sink_mass)
logger.info('Sink mass goes below %s; excess mass is now %s',
minimum_sink_mass.in_(units.MSun),
excess_star_mass.in_(units.MSun))
s.mass = minimum_sink_mass
else:
s.mass -= total_star_mass_nearby
else:
excess_star_mass += total_star_mass_nearby
logger.info(
'Sink mass is already <= minimum mass allowed; '
'excess mass is now %s', excess_star_mass.in_(units.MSun))
logger.info('Sink mass after reduction: %s', s.mass.in_(units.MSun))
# Reduce all sinks in group equally with the excess star mass
logger.info('Reducing all sink mass equally with excess star mass...')
mass_ratio = 1 - excess_star_mass / group.total_mass()
group.mass *= mass_ratio
logger.info("Total sink mass in group: %s",
group.total_mass().in_(units.MSun))
if shrink_sinks:
group.radius = ((group.mass / group.initial_density) /
(4 / 3 * numpy.pi))**(1 / 3)
logger.info("New radii: %s", group.radius.in_(units.pc))
return new_stars
def rgb_frame(
stars,
dryrun=False,
vmax=None,
percentile=0.9995,
multi_psf=False,
sourcebands="ubvri",
image_width=12. | units.parsec,
image_size=[1024, 1024],
mapper_factory=None,
gas=None,
mapper_code=None,
zoom_factor=1.0,
psf_type="hubble",
psf_sigma=1.0,
verbose=True,
):
if gas is None:
gas = Particles()
if verbose:
print("luminosities..")
for band in sourcebands:
setattr(
stars,
band + "_band",
4 * numpy.pi * stars.radius**2 * filter_band_flux(
"bess-" + band + ".pass",
lambda x: B_lambda(x, stars.temperature),
),
)
if verbose:
print("..raw images..")
if mapper_code != "gridify":
# Use mapper to make raw (pre-convolved) images
mapper = mapper_factory()
stars_in_mapper = mapper.particles.add_particles(stars)
gas_in_mapper = mapper.particles.add_particles(gas)
mapper.parameters.projection_direction = [0, 0, 1]
mapper.parameters.upvector = [0, -1, 0]
raw_images = dict()
for band in sourcebands:
assign_weights_and_opacities(
band,
stars_in_mapper,
gas_in_mapper,
stars,
gas,
Nstar=500,
)
# mapper.particles.weight = getattr(
# stars,
# band+"_band"
# ).value_in(units.LSun)
im = mapper.image.pixel_value
raw_images[band] = numpy.fliplr(im)
mapper.stop()
else:
# Use simpler python mapping script
from .gridify import map_to_grid
stars_in_mapper = stars.copy()
gas_in_mapper = gas.copy()
raw_images = dict()
for band in sourcebands:
assign_weights_and_opacities(
band,
stars_in_mapper,
gas_in_mapper,
stars,
gas,
Nstar=500,
)
allparticles = Particles()
allparticles.add_particles(stars_in_mapper)
allparticles.add_particles(gas_in_mapper)
im = map_to_grid(
allparticles.x,
allparticles.y,
weights=allparticles.weight,
image_size=image_size,
image_width=image_width,
)
raw_images[band] = im
convolved_images = dict()
if verbose:
print("..convolving..")
if psf_type == "hubble":
psf = get_psf(zoom_factor=zoom_factor)
if multi_psf:
a = numpy.arange(image_size[0]) / float(image_size[0] - 1)
b = numpy.arange(image_size[1]) / float(image_size[1] - 1)
w1 = numpy.outer(a, b)
w2 = numpy.outer(1. - a, b)
w3 = numpy.outer(a, 1. - b)
w4 = numpy.outer(1. - a, 1. - b)
for key, val in list(raw_images.items()):
im1 = Convolve(val, psf[key + '0'])
im2 = Convolve(val, psf[key + '1'])
im3 = Convolve(val, psf[key + '2'])
im4 = Convolve(val, psf[key + '3'])
convolved_images[key] = (w1 * im1 + w2 * im2 + w3 * im3 +
w4 * im4)
else:
for key, val in list(raw_images.items()):
im1 = Convolve(val, psf[key + '0'])
convolved_images[key] = im1
elif psf_type == "gaussian":
for key, val in list(raw_images.items()):
im1 = gaussian_filter(val, sigma=psf_sigma, order=0)
convolved_images[key] = im1
if verbose:
print("..conversion to rgb")
filter_data = get_filter_data()
source = [filter_data['bess-' + x + '.pass'] for x in sourcebands]
target = [xyz_data['x'], xyz_data['y'], xyz_data['z']]
conv = ColorConverter(source, target)
ubv = numpy.array([convolved_images[x] for x in sourcebands])
xyz = numpy.tensordot(conv.conversion_matrix, ubv, axes=(1, 0))
conv_xyz_to_lin = XYZ_to_sRGB_linear()
srgb_l = numpy.tensordot(
conv_xyz_to_lin.conversion_matrix,
xyz,
axes=(1, 0),
)
if dryrun or vmax is None:
flat_sorted = numpy.sort(srgb_l.flatten())
n = len(flat_sorted)
vmax = flat_sorted[int(1. - 3 * (1. - percentile) * n)]
print(("vmax:", vmax))
if dryrun:
return vmax
conv_lin_to_sRGB = sRGB_linear_to_sRGB()
srgb = conv_lin_to_sRGB.convert(srgb_l / vmax)
# r = numpy.fliplr(srgb[0, :, :])
# g = numpy.fliplr(srgb[1, :, :])
# b = numpy.fliplr(srgb[2, :, :])
rgb = numpy.zeros(image_size[0] * image_size[1] * 3)
rgb = rgb.reshape(image_size[0], image_size[1], 3)
rgb[:, :, 0] = srgb[0, :, :].T
rgb[:, :, 1] = srgb[1, :, :].T
rgb[:, :, 2] = srgb[2, :, :].T
image = dict(
pixels=rgb,
size=rgb.size,
)
return vmax, image
class MinimalWorkingExample(object):
def __init__(self,
total_N=16000, #total number of disc particles
tend=250. | units.yr, #End time of the simulation
Mstar=1. | units.MSun, #Mass of the accreting star
dt = 1. | units.yr, #timestep
temp = 25. | units.K, #temperature of the ISM
v_ism = 3.0 | units.kms, #Velocity of the ISM with respect to the star
n_core = 4, #Number of cores used by the community codes
mu = 2.3, #mean molecular weight
theta = 0., #angle of the disc with respect to the flow, 0 means disc rotation axis is parallel with flow direction
n_dens = 5e6 | (units.cm)**(-3.), #Number density of the ism
dirname = './'):
self.disc_N = total_N
self.tend = tend
self.Mstar = Mstar
self.dt = dt
self.mu = mu * constants.proton_mass
self.T = temp
self.cs = ((temp*constants.kB)/self.mu).sqrt()
# print "cs=", self.cs.in_(units.kms)
self.v_ism = v_ism
self.ism_n_dens = n_dens #number density of the ism
self.ism_dens = self.mu*self.ism_n_dens #mass density of the ism
self.converter = nbody_system.nbody_to_si(self.Mstar, 1. | units.AU)
self.disc_angle = numpy.radians(theta)
self.n_core = n_core
self.dirname = dirname
self.discfraction = 0.01
self.disc_Rmin = self.converter.to_generic(10.| units.AU).value_in(nbody_system.length)
self.disc_Rmax = self.converter.to_generic(100.| units.AU).value_in(nbody_system.length)
self.filename = "DiskWind.h5"
print 'Files will be saved in ', self.dirname
def initialize_star(self):
self.star=Particles(1)
self.star.mass=self.Mstar
self.star.radius= 1. | units.AU
self.star.x=0.|units.AU
self.star.y=0.|units.AU
self.star.z=0.|units.AU
self.star.vx=0.|units.kms
self.star.vy=0.|units.kms
self.star.vz=0.|units.kms
def initialize_data(self):
self.cylinder_radius = 500. | units.AU
self.cylinder_length = 2.*self.cylinder_radius
self.cylinder_vol = self.cylinder_length*numpy.pi*self.cylinder_radius**2.
self.sph_particle_dens = self.ism_dens/(0.01*self.Mstar/self.disc_N)
self.initialize_star()
self.initialize_ism()
def initialize_ism(self):
#specific internal energy of the ism, i.e. internal energy per unit mass
#In case of an isothermal EOS, Gadget requires the sound-speed squared as input parameter instead of the thermal energy per unit mass
self.ism_u = self.cs**2.
self.ism_slice_length = self.v_ism*self.dt
self.ism_slice_vol = self.ism_slice_length*numpy.pi*self.cylinder_radius**2.0
self.ism_slice_mass = self.ism_dens*self.ism_slice_vol
self.ism_mass = self.ism_dens*self.cylinder_vol
#The x-coordinate of the inflow
self.offset = self.cylinder_radius
self.ism_slice_N = int(round(self.sph_particle_dens*self.ism_slice_vol))
self.total_N = self.ism_slice_N*(self.cylinder_length/self.ism_slice_length)
def make_slice(self):
new_slice = Particles(self.ism_slice_N)
rho = numpy.sqrt(numpy.random.uniform(0,1,self.ism_slice_N))*self.cylinder_radius.value_in(units.AU)
phi = numpy.random.uniform(0,2.0*numpy.pi, self.ism_slice_N)
new_slice.x = numpy.random.uniform(self.ism_slice_length.value_in(units.AU), 0, self.ism_slice_N) - self.offset.value_in(units.AU) | units.AU
new_slice.y = rho*numpy.sin(phi) | units.AU
new_slice.z = rho*numpy.cos(phi) | units.AU
new_slice.vx = 1000*self.v_ism
new_slice.vy = 0.0 | units.kms
new_slice.vz = 0.0 | units.kms
new_slice.mass = self.ism_slice_mass/self.ism_slice_N
new_slice.u = self.ism_u
return new_slice
def create_disc(self):
#The following line makes sure masses for ISM and disc particles are equal:
self.discfraction = (self.disc_N*(self.ism_slice_mass/self.ism_slice_N))/self.Mstar
T_disc = self.T
cs_disc = ((T_disc*constants.kB)/self.mu).sqrt()
densitypower = 1.5
g2=2-densitypower
k_out=((1+self.discfraction)/self.disc_Rmax**3)**0.5
sigma_out=g2*self.discfraction/(2*numpy.pi*self.disc_Rmax**densitypower*(self.disc_Rmax**g2-self.disc_Rmin**g2))
q_out = self.converter.to_generic(cs_disc).value_in(nbody_system.length/nbody_system.time)/(numpy.pi*sigma_out/k_out)
print "Number of disk particles:", self.disc_N
proto=ProtoPlanetaryDisk(self.disc_N,
convert_nbody=self.converter,
discfraction=self.discfraction,
densitypower=1.5,
thermalpower=0,
Rmin=self.disc_Rmin,
Rmax=self.disc_Rmax,
q_out=q_out)
disc=proto.result
print "The mass of a disc particle = ", disc.mass[0].value_in(units.kg)
#Rotate 90 degrees with respect to the z-axis and then theta degrees with respect to the y-axis
temp_x = disc[:].x
temp_y = disc[:].y
temp_z = disc[:].z
temp_vx = disc[:].vx
temp_vy = disc[:].vy
temp_vz = disc[:].vz
disc.x = temp_z*numpy.cos(self.disc_angle) - temp_y*numpy.sin(self.disc_angle)
disc.y = temp_z*numpy.sin(self.disc_angle) + temp_y*numpy.cos(self.disc_angle)
disc.z = -temp_x
disc.vx = temp_vz*numpy.cos(self.disc_angle) - temp_vy*numpy.sin(self.disc_angle)
disc.vy = temp_vz*numpy.sin(self.disc_angle) + temp_vy*numpy.cos(self.disc_angle)
disc.vz = -temp_vx
return disc
def evolve_model(self):
self.initialize_data()
self.ism_code = Gadget2(self.converter, number_of_workers=self.n_core)#, debugger='gdb')
self.ism_code.parameters.time_max = 1024*self.dt
self.ism_code.parameters.n_smooth = 64
self.ism_code.parameters.n_smooth_tol = 2./64.
self.ism_code.parameters.artificial_viscosity_alpha = 0.1
self.ism_code.parameters.epsilon_squared = (1. | units.AU)**2.
self.all_particles = Particles()
write_set_to_file(self.star, self.filename, "hdf5", append_to_file=False)
write_set_to_file(self.all_particles, self.filename, "hdf5", append_to_file=False)
self.initial_disc_particles = self.create_disc()
self.all_particles.add_particles(self.initial_disc_particles)
self.ism_code.gas_particles.add_particles(self.initial_disc_particles)
#You can only add a sink after adding gas particles
#starinsph refers to the corresponding particle set/id in the community code
starinsph = self.ism_code.dm_particles.add_particles(self.star)
#Use the build-in sink particle routine from amuse.ext.sink. Sink_radius needs to be defined manually otherwise the particle radius in gadget is taken,
#which does not corresponding to the particle radius in the framework (since 'particles' in gadget do not have a radius, it is set to 0.01 | generic_unit_system.length
#and corresponds to the gas gravitational smoothing epsilon.
sink = new_sink_particles(starinsph, sink_radius= self.star.radius)
self.channel_from_ismcode_to_framework = self.ism_code.gas_particles.new_channel_to(self.all_particles)
time = 0. | units.yr
while time <= (self.tend+self.dt/2.):
print "Adding new slice of ISM..."
newslice=self.make_slice()
self.ism_code.gas_particles.add_particles(newslice)
self.all_particles.add_particles(newslice)
start = timing.time()
print "======================================================="
print "Evolving to time = ", time.value_in(units.yr), " of ", self.tend.value_in(units.yr)," years..."
self.ism_code.evolve_model(time)
print "This took ", (timing.time() - start), " s"
out_of_bounds = self.ism_code.gas_particles.select_array(lambda x,y,z:(x > self.cylinder_radius)|((z**2+y**2).sqrt() >= self.cylinder_radius), ["x","y","z"])
if len(out_of_bounds)>0:
print "Removing ", len(out_of_bounds), " particles from the code because they were out of bounds"
self.ism_code.gas_particles.remove_particles(out_of_bounds)
self.ism_code.gas_particles.synchronize_to(self.all_particles)
sink.accrete(self.ism_code.gas_particles)
write_set_to_file(self.star, self.filename, "hdf5")
write_set_to_file(self.all_particles, self.filename, "hdf5")
time += self.dt
print "======================================================="
self.ism_code.stop()
def main():
"Make a star cluster"
set_printing_strategy(
"custom",
preferred_units=[units.MSun, units.parsec, units.yr, units.kms],
precision=5,
)
clustertemplate = "TESTCluster_%08i"
args = new_argument_parser()
sinks = args.sinks
cluster_model_number = args.cluster_model_number
star_distribution = args.star_distribution
# gas_distribution = args.gas_distribution
king_parameter_w0 = args.king_parameter_w0
fractal_parameter_fd = args.fractal_parameter_fd
initial_mass_function = args.initial_mass_function.lower()
number_of_stars = args.number_of_stars
if args.cluster_mass != 0:
cluster_mass = args.cluster_mass | units.MSun
else:
cluster_mass = False
upper_mass_limit = args.upper_mass_limit | units.MSun
lower_mass_limit = args.lower_mass_limit | units.MSun
effective_radius = args.effective_radius | units.parsec
metallicity = args.metallicity
# virial_ratio = args.virial_ratio
filetype = args.filetype
# not implemented yet
# initial_binary_fraction = args.initial_binary_fraction
np.random.seed(cluster_model_number)
if not (number_of_stars or cluster_mass or sinks):
print(
"no number of stars, cluster mass or origin sinks given, exiting"
)
exit()
if sinks is not None:
sinks = read_set_from_file(sinks, "amuse")
stars = Particles()
for sink in sinks:
try:
velocity_dispersion = sink.u.sqrt()
except AttributeError:
velocity_dispersion = args.velocity_dispersion | units.kms
new_stars = new_stars_from_sink(
sink,
upper_mass_limit=upper_mass_limit,
lower_mass_limit=lower_mass_limit,
default_radius=effective_radius,
velocity_dispersion=velocity_dispersion,
initial_mass_function=initial_mass_function,
# logger=logger,
)
stars.add_particles(
new_stars
)
else:
stars = new_star_cluster(
stellar_mass=cluster_mass,
initial_mass_function=initial_mass_function,
upper_mass_limit=upper_mass_limit,
lower_mass_limit=lower_mass_limit,
number_of_stars=number_of_stars,
effective_radius=effective_radius,
star_distribution=star_distribution,
star_distribution_w0=king_parameter_w0,
star_distribution_fd=fractal_parameter_fd,
star_metallicity=metallicity,
)
print(
"%i stars generated (%s)"
% (len(stars), stars.total_mass().in_(units.MSun))
)
if args.clustername != "auto":
clustertemplate = args.clustername + "%s"
stars_file_exists = True
sinks_file_exists = True
N = -1
while (stars_file_exists or sinks_file_exists):
N += 1
starsfilename = (
clustertemplate % N
+ "-stars." + filetype
)
stars_file_exists = os.path.isfile(starsfilename)
sinksfilename = (
clustertemplate % N
+ "-sinks." + filetype
)
sinks_file_exists = os.path.isfile(sinksfilename)
write_set_to_file(stars, starsfilename, filetype)
if sinks is not None:
write_set_to_file(sinks, sinksfilename, filetype)
def main():
# Fixed settings
stellar_evolution = True
se_code = "SeBa"
length_unit = units.parsec
dpi = 600
percentile = 0.9995 # for determining vmax
# Parse arguments
args = new_argument_parser()
starsfilename = args.starsfilename
gasfilename = args.gasfilename
followfilename = args.followfilename
imagefilename = args.imagefilename
imagetype = args.imagetype
vmax = args.vmax if args.vmax > 0 else None
n_fieldstars = args.n_fieldstars
filetype = args.filetype
contours = args.contours
np.random.seed(args.seed)
plot_axes = args.plot_axes
angle_x = args.angle_x | units.deg
angle_y = args.angle_y | units.deg
angle_z = args.angle_z | units.deg
sourcebands = args.sourcebands
psf_type = args.psf_type.lower()
psf_sigma = args.psf_sigma
age = args.age | units.Myr
image_width = args.width | units.parsec
pixels = args.pixels
frames = args.frames
if followfilename is not None:
use_com = True
else:
use_com = args.use_com
x_offset = args.x_offset | units.parsec
y_offset = args.y_offset | units.parsec
z_offset = args.z_offset | units.parsec
extinction = args.calculate_extinction
# Derived settings
if psf_type not in ["hubble", "gaussian"]:
print(("Invalid PSF type: %s" % psf_type))
exit()
image_size = [pixels, pixels]
# If the nr of pixels is changed, zoom the PSF accordingly.
zoom_factor = pixels / 2048.
if starsfilename:
stars = read_set_from_file(
starsfilename,
filetype,
close_file=True,
)
if stellar_evolution and (age > 0 | units.Myr):
print((
"Calculating luminosity/temperature for %s old stars..."
% (age)
))
evolve_to_age(stars, age, stellar_evolution=se_code)
if use_com:
if followfilename is not None:
followstars = read_set_from_file(
followfilename, filetype, close_file=True,
)
center_on_these_stars = followstars.get_intersecting_subset_in(
stars,
)
else:
center_on_these_stars = stars
com = center_on_these_stars.center_of_mass()
x_offset, y_offset, z_offset = com
stars.x -= x_offset
stars.y -= y_offset
stars.z -= z_offset
else:
stars = Particles()
if n_fieldstars:
minage = 400 | units.Myr
maxage = 12 | units.Gyr
fieldstars = new_field_stars(
n_fieldstars,
width=image_width,
height=image_width,
)
fieldstars.age = (
minage
+ (
np.random.sample(n_fieldstars)
* (maxage - minage)
)
)
evolve_to_age(fieldstars, 0 | units.yr, stellar_evolution=se_code)
stars.add_particles(fieldstars)
if gasfilename:
gas = read_set_from_file(
gasfilename,
filetype,
close_file=True,
)
if use_com:
if stars.is_empty():
com = gas.center_of_mass()
x_offset, y_offset, z_offset = com
gas.x -= x_offset
gas.y -= y_offset
gas.z -= z_offset
# Gadget and Fi disagree on the definition of h_smooth.
# For gadget, need to divide by 2 to get the Fi value (??)
gas.h_smooth *= 0.5
gas.radius = gas.h_smooth
# Select only the relevant gas particles (plus a margin)
minx = (1.1 * -image_width/2)
maxx = (1.1 * image_width/2)
miny = (1.1 * -image_width/2)
maxy = (1.1 * image_width/2)
gas_ = gas.select(
lambda x, y:
x > minx
and x < maxx
and y > miny
and y < maxy,
["x", "y"]
)
gas = gas_
else:
gas = Particles()
# gas.h_smooth = 0.05 | units.parsec
converter = nbody_system.nbody_to_si(
stars.total_mass() if not stars.is_empty() else gas.total_mass(),
image_width,
)
# Initialise figure and axes
fig = initialise_image(
dpi=dpi,
image_size=image_size,
length_unit=length_unit,
image_width=image_width,
plot_axes=plot_axes,
x_offset=x_offset,
y_offset=y_offset,
z_offset=z_offset,
)
ax = fig.get_axes()[0]
xmin, xmax = ax.get_xlim()
ymin, ymax = ax.get_ylim()
for frame in range(frames):
fig = initialise_image(fig)
if (frame != 0) or (frames == 1):
if not stars.is_empty():
rotate(stars, angle_x, angle_y, angle_z)
if not gas.is_empty():
rotate(gas, angle_x, angle_y, angle_z)
image, vmax = make_image(
stars=stars if not stars.is_empty() else None,
gas=gas if not gas.is_empty() else None,
converter=converter,
image_width=image_width,
image_size=image_size,
percentile=percentile,
calc_temperature=True,
age=age,
vmax=vmax,
sourcebands=sourcebands,
zoom_factor=zoom_factor,
psf_type=psf_type,
psf_sigma=psf_sigma,
return_vmax=True,
extinction=extinction,
)
if not stars.is_empty():
ax.imshow(
image,
origin='lower',
extent=[
xmin,
xmax,
ymin,
ymax,
],
)
if contours and not gas.is_empty():
gascontours = column_density_map(
gas,
zoom_factor=zoom_factor,
image_width=image_width,
image_size=image_size,
)
gascontours[np.isnan(gascontours)] = 0.0
vmax = np.max(gascontours) / 2
# vmin = np.min(image[np.where(image > 0.0)])
vmin = vmax / 100
levels = 10**(
np.linspace(
np.log10(vmin),
np.log10(vmax),
num=5,
)
)[1:]
# print(vmin, vmax)
# print(levels)
ax.contour(
origin='lower',
levels=levels,
colors="white",
linewidths=0.1,
extent=[
xmin,
xmax,
ymin,
ymax,
],
)
else:
image = column_density_map(
gas,
image_width=image_width,
image_size=image_size,
)
ax.imshow(
image,
origin='lower',
extent=[
xmin,
xmax,
ymin,
ymax,
],
cmap="gray",
)
if frames > 1:
savefilename = "%s-%06i.%s" % (
imagefilename if imagefilename is not None else "test",
frame,
imagetype,
)
else:
savefilename = "%s.%s" % (
imagefilename if imagefilename is not None else "test",
imagetype,
)
plt.savefig(
savefilename,
dpi=dpi,
)
def main():
mode = []
# Fixed settings
stellar_evolution = True
se_code = "SeBa"
length_unit = units.parsec
dpi = 600
percentile = 0.9995 # for determining vmax
# Parse arguments
args = new_argument_parser()
starsfilename = args.starsfilename
gasfilename = args.gasfilename
imagefilename = args.imagefilename
imagetype = args.imagetype
vmax = args.vmax if args.vmax > 0 else None
n_fieldstars = args.n_fieldstars
filetype = args.filetype
contours = args.contours
np.random.seed(args.seed)
plot_axes = args.plot_axes
angle_x = args.angle_x | units.deg
angle_y = args.angle_y | units.deg
angle_z = args.angle_z | units.deg
sourcebands = args.sourcebands
psf_type = args.psf_type.lower()
psf_sigma = args.psf_sigma
age = args.age | units.Myr
image_width = args.width | units.parsec
pixels = args.pixels
frames = args.frames
# Derived settings
if args.calculate_extinction:
mode.append("extinction")
if psf_type not in ["hubble", "gaussian"]:
print(("Invalid PSF type: %s" % psf_type))
exit()
image_size = [pixels, pixels]
# If the nr of pixels is changed, zoom the PSF accordingly.
zoom_factor = pixels / 2048.
if starsfilename:
stars = read_set_from_file(
starsfilename,
filetype,
close_file=True,
)
if stellar_evolution and (age > 0 | units.Myr):
print(("Calculating luminosity/temperature for %s old stars..." %
(age)))
evolve_to_age(stars, age, stellar_evolution=se_code)
com = stars.center_of_mass()
stars.position -= com
else:
stars = Particles()
if n_fieldstars:
minage = 400 | units.Myr
maxage = 12 | units.Gyr
fieldstars = new_field_stars(
n_fieldstars,
width=image_width,
height=image_width,
)
fieldstars.age = (minage + (np.random.sample(n_fieldstars) *
(maxage - minage)))
evolve_to_age(fieldstars, 0 | units.yr, stellar_evolution=se_code)
stars.add_particles(fieldstars)
if not stars.is_empty():
mode.append("stars")
if gasfilename:
gas = read_set_from_file(
gasfilename,
filetype,
close_file=True,
)
if "stars" not in mode:
com = gas.center_of_mass()
gas.position -= com
# Gadget and Fi disagree on the definition of h_smooth.
# For gadget, need to divide by 2 to get the Fi value (??)
gas.h_smooth *= 0.5
gas.radius = gas.h_smooth
else:
gas = Particles()
if not gas.is_empty():
mode.append("gas")
if contours:
mode.append("contours")
# gas.h_smooth = 0.05 | units.parsec
converter = nbody_system.nbody_to_si(
stars.total_mass() if "stars" in mode else gas.total_mass(),
image_width,
)
# Initialise figure and axes
fig = initialise_image(
dpi=dpi,
image_size=image_size,
length_unit=length_unit,
image_width=image_width,
plot_axes=plot_axes,
)
ax = fig.get_axes()[0]
xmin, xmax = ax.get_xlim()
ymin, ymax = ax.get_ylim()
for frame in range(frames):
fig = initialise_image(fig)
if (frame != 0) or (frames == 1):
if not stars.is_empty():
rotate(stars, angle_x, angle_y, angle_z)
if not gas.is_empty():
rotate(gas, angle_x, angle_y, angle_z)
image, vmax = make_image(
stars,
gas,
mode=mode,
converter=converter,
image_width=image_width,
image_size=image_size,
percentile=percentile,
calc_temperature=True,
age=age,
vmax=vmax,
sourcebands=sourcebands,
zoom_factor=zoom_factor,
psf_type=psf_type,
psf_sigma=psf_sigma,
return_vmax=True,
)
if "stars" in mode:
ax.imshow(
image,
origin='lower',
extent=[
xmin,
xmax,
ymin,
ymax,
],
)
if ("contours" in mode) and ("gas" in mode):
gascontours = column_density_map(
gas,
zoom_factor=zoom_factor,
image_width=image_width,
image_size=image_size,
)
gascontours[np.isnan(gascontours)] = 0.0
vmax = np.max(gascontours) / 2
# vmin = np.min(image[np.where(image > 0.0)])
vmin = vmax / 100
levels = 10**(np.linspace(
np.log10(vmin),
np.log10(vmax),
num=5,
))[1:]
# print(vmin, vmax)
# print(levels)
ax.contour(
gascontours,
origin='lower',
levels=levels,
colors="white",
linewidths=0.1,
extent=[
xmin,
xmax,
ymin,
ymax,
],
)
else:
image = column_density_map(
gas,
image_width=image_width,
image_size=image_size,
)
ax.imshow(
image,
origin='lower',
extent=[
xmin,
xmax,
ymin,
ymax,
],
cmap="gray",
)
plt.savefig(
"%s-%06i.%s" % (
imagefilename,
frame,
imagetype,
),
dpi=dpi,
)
class StellarDynamicsCode:
"""Wraps around stellar dynamics code, supports collisions"""
def __init__(
self,
converter=None,
# star_code=ph4,
star_code=Petar,
# star_code=Hermite,
logger=None,
handle_stopping_conditions=False,
# mode="cpu",
time_offset=0 | nbody_system.time,
stop_after_each_step=False,
number_of_workers=8,
settings=None,
**kwargs):
self.__name__ = "StellarDynamics"
self.logger = logger or logging.getLogger(__name__)
if settings is None:
from ekster_settings import settings
print("WARNING: using default settings!")
logger.info("WARNING: using default settings!")
self.settings = settings
epsilon_squared = settings.epsilon_stars**2
self.typestr = "Nbody"
self.star_code = star_code
try:
self.namestr = self.star_code.__name__
except AttributeError:
self.namestr = "unknown name"
self.handle_stopping_conditions = \
handle_stopping_conditions
self.__current_state = "stopped"
self.__state = {}
self.__particles = Particles()
self.__stop_after_each_step = (stop_after_each_step if self.star_code
is not Petar else False)
if converter is not None:
self.unit_converter = converter
else:
self.unit_converter = nbody_system.nbody_to_si(
settings.star_mscale,
settings.star_rscale,
)
# TODO: modify to allow N-body units
if time_offset is None:
time_offset = 0. | units.Myr
self.__time_offset = self.unit_converter.to_si(time_offset)
self.code = self.new_code(converter=self.unit_converter,
star_code=star_code,
epsilon_squared=epsilon_squared,
number_of_workers=number_of_workers,
**kwargs)
self.parameters_to_default(star_code=star_code, )
if self.__stop_after_each_step:
# self.code.commit_particles()
self.stop(save_state=True)
# print("Stopped/saved")
else:
self.save_state()
def new_code(
self,
converter=None,
star_code=Hermite,
redirection="null",
mode="cpu",
number_of_workers=8,
# handle_stopping_conditions=False,
**kwargs):
if hasattr(available_codes, star_code):
star_code = getattr(available_codes, star_code)
if star_code is ph4:
code = star_code(converter,
mode=mode,
redirection=redirection,
number_of_workers=number_of_workers,
**kwargs)
elif star_code is Hermite:
code = star_code(
converter,
number_of_workers=number_of_workers,
redirection=redirection,
)
elif star_code is PhiGRAPE:
code = star_code(
converter,
number_of_workers=number_of_workers,
redirection=redirection,
)
elif star_code is BHTree:
code = star_code(
converter,
redirection=redirection,
)
elif star_code is Pentacle:
code = star_code(
converter,
# redirection=redirection,
redirection="none",
)
elif star_code is Petar:
code = star_code(
converter,
mode=mode,
# redirection=redirection,
redirection="none",
# number_of_workers=number_of_workers,
**kwargs)
else:
raise Exception("Code not found: %s" % star_code)
self.__current_state = "started"
return code
def parameters_to_default(
self,
star_code=Hermite,
):
"Set default parameters"
settings = self.settings
epsilon_squared = settings.epsilon_stars**2
logger = self.logger
param = self.code.parameters
param.epsilon_squared = epsilon_squared
if star_code is ph4:
# Set the parameters explicitly to some default
# param.block_steps = False
# Force ph4 to synchronise to the exact time requested - important
# for Bridge!
param.force_sync = True
# param.gpu_id = something
param.initial_timestep_fac = 0.0625
param.initial_timestep_limit = 0.03125
# param.initial_timestep_median = 8.0
# param.manage_encounters = 4
# # We won't use these stopping conditions anyway
# param.stopping_condition_maximum_density = some HUGE number
# param.stopping_condition_maximum_internal_energy = inf
# param.stopping_condition_minimum_density = - huge
# param.stopping_condition_minimum_internal_energy = - big number
# param.stopping_conditions_number_of_steps = 1
# param.stopping_conditions_out_of_box_size = 0 | units.m
# param.stopping_conditions_out_of_box_use_center_of_mass = True
# param.stopping_conditions_timeout = 4.0 | units.s
# param.sync_time = 0.0 | units.s
# param.timestep_parameter = 0.0
# param.total_steps = False
# param.use_gpu = False
# param.zero_step_mode = False
elif star_code is Hermite:
# Force Hermite to sync to the exact time requested - see
# force_sync for ph4
param.end_time_accuracy_factor = 0
elif star_code is Petar:
# Set the parameters explicitly to some default
param.theta = settings.stellar_dynamics_theta
logger.info("Old r_out value: %s", param.r_out.in_(units.pc))
param.r_out = settings.stellar_dynamics_r_out
param.ratio_r_cut = settings.stellar_dynamics_ratio_r_cut
logger.info("Old r_bin value: %s", param.r_bin.in_(units.pc))
param.r_bin = settings.stellar_dynamics_r_bin
# param.r_search_min = 0 | units.pc
# very small = technically disabled
param.r_search_min = settings.stellar_dynamics_r_search_min
param.dt_soft = settings.stellar_dynamics_dt_soft
# param.dt_soft = self.unit_converter.to_si(
# 2**-8 | nbody_system.time
# )
# settings.timestep_bridge / 4 # 0 | units.Myr
# param.r_out = 10 * settings.epsilon_stars
# dt_soft: 9.765625e-06 Myr default: 0.0 Myr
# epsilon_squared: 0.0001 parsec**2 default: 0.0 parsec**2
# r_bin: 0.000137167681417 parsec default: 0.0 parsec
# r_out: 0.00171459601771 parsec default: 0.0 parsec
# r_search_min: 0.00206443388608 parsec default: 0.0 parsec
# ratio_r_cut: 0.1 default: 0.1
# r_in = 0.00043686
# r_out = 0.0043686
# r_bin = 0.00034949
# r_search_min = 0.0056792
# vel_disp = 0.89469
# dt_soft = 0.00048828
def evolve_model(self, end_time):
"""
Evolve model, handle collisions when they occur
"""
if self.__stop_after_each_step:
# print("Code will be stopped after each step")
if self.__current_state == "stopped":
# print("Code is currently stopped - restarting")
self.restart()
result = 0
time_unit = end_time.unit
time_fraction = 1 | units.s
print("START model time: %s -> end_time: %s" % (
self.model_time.in_(units.Myr),
end_time.in_(units.Myr),
))
self.logger.info(
"Starting evolve of %s, model time is %s, end time is %s",
self.__name__,
self.model_time.in_(time_unit),
end_time.in_(time_unit),
)
while self.model_time < end_time:
print("%s < %s, continuing" % (
self.model_time.in_(time_unit),
end_time.in_(time_unit),
))
if self.model_time >= (end_time - time_fraction):
print("but %s >= (%s-%s), not continuing" %
(self.model_time.in_(time_unit), end_time.in_(time_unit),
time_fraction.in_(time_unit)))
break
if not self.code.particles.is_empty():
print("Starting evolve_model of stellar_dynamics")
result = self.code.evolve_model(end_time - self.__time_offset)
print("Finished evolve_model of stellar_dynamics")
else:
self.logger.info(
"No particles, skipping evolve and readjusting time offset"
)
print(
"Skipping evolve_model of stellar_dynamics, no particles!")
self.__time_offset = end_time
result = 0
if self.__stop_after_each_step:
# print("Now stopping code")
self.stop(save_state=True)
print("FINISH model time: %s > end_time: %s" % (
self.model_time.in_(units.Myr),
end_time.in_(units.Myr),
))
self.logger.info(
"Finishing evolve of %s, model time is %s, end time is %s",
self.__name__,
self.model_time.in_(time_unit),
end_time.in_(time_unit),
)
return result
@property
def model_time(self):
"""Return code model_time"""
if self.__current_state != "stopped":
time = self.code.model_time + self.__time_offset
return time
time = self.__last_time
return time
@property
def particles(self):
"""Return particles"""
if self.__stop_after_each_step:
return self.__particles
# if self.__current_state is not "stopped":
# return self.code.particles
else:
return self.code.particles
@property
def parameters(self):
"""Return code parameters"""
if self.__current_state != "stopped":
parameters = self.code.parameters
else:
parameters = self.__state["parameters"]
return parameters
# TODO: make sure this parameter set is synchronised with code.parameters
# def parameters(self):
# """Return code parameters"""
# self.__parameters = self.code.parameters.copy()
# return self.__parameters
@property
def stopping_conditions(self):
"""Return stopping conditions for dynamics code"""
return self.code.stopping_conditions
@property
def commit_particles(self):
return self.code.commit_particles
def get_gravity_at_point(self, *list_arguments, **keyword_arguments):
"""Return gravity at specified point"""
return self.code.get_gravity_at_point(*list_arguments,
**keyword_arguments)
def get_potential_at_point(self, *list_arguments, **keyword_arguments):
"""Return potential at specified point"""
return self.code.get_potential_at_point(*list_arguments,
**keyword_arguments)
def save_state(self):
"""
Store current settings
"""
self.__state["parameters"] = self.code.parameters.copy()
self.__state["converter"] = self.unit_converter
self.__state["star_code"] = self.star_code
self.__state["model_time"] = self.code.model_time
self.__state["redirection"] = "null" # FIXME
self.__state["mode"] = "cpu" # FIXME
self.__state["handle_stopping_conditions"] = \
self.handle_stopping_conditions
self.__last_time = self.model_time
def save_particles(self):
"""
Store the current particleset, but keep the same particleset!
"""
self.__particles.remove_particles(self.__particles)
self.__particles.add_particles(self.code.particles)
def stop_and_restart(self):
"""
Store current settings and restart gravity code from saved state
"""
self.stop(save_state=True)
self.restart()
def restart(self):
"""
Restart gravity code from saved state
"""
# print("Restarting")
self.code = self.new_code(
converter=self.__state["converter"],
star_code=self.__state["star_code"],
redirection=self.__state["redirection"],
mode=self.__state["mode"],
handle_stopping_conditions=self.
__state["handle_stopping_conditions"],
)
self.code.particles.add_particles(self.__particles)
print(self.__state["parameters"])
if self.star_code is Petar:
for name in self.__state["parameters"].names():
if name != "timestep":
setattr(self.code.parameters, name,
getattr(self.__state["parameters"], name))
else:
self.code.parameters.reset_from_memento(self.__state["parameters"])
self.__current_state = "restarted"
def stop(self, save_state=True, **keyword_arguments):
"""Stop code"""
if save_state:
self.save_state(**keyword_arguments)
self.save_particles(**keyword_arguments)
stopcode = self.code.stop(**keyword_arguments)
self.__current_state = "stopped"
return stopcode