def get_ps_from_yt(pltfile, particle_type="all"): """Take a FLASH plotfile and return an AMUSE particle set. Keyword arguments: pltfile -- The FLASH hdf5 plot or checkpoint file. Note a plotfile means there must also be a particle file with the same ID number. particle_type -- the type of particle that you want from FLASH. Can be: star -- active FLASH particle type sink -- sink FLASH particle type any other string or none returns all particle types. The default is all Returns: stars -- An AMUSE particle set. """ import yt ds = yt.load(pltfile) dd = ds.all_data() all_parts_mass = dd['particle_mass'].in_units('Msun') all_parts_pos = dd['particle_position'].in_units('cm') all_parts_vel = dd['particle_velocity'].in_units('cm/s') if (particle_type == 'star' or particle_type == 'stars'): # Index for only the star particles, and not including sink particles. mass_ind = np.where(dd['particle_type'].v == 1) elif (particle_type == 'sink' or particle_type == 'sinks'): mass_ind = np.where(dd['particle_type'].v == 2) else: mass_ind = np.arange(len(all_parts_mass)) all_parts_mass = all_parts_mass[mass_ind] all_parts_pos = all_parts_pos[mass_ind] all_parts_vel = all_parts_vel[mass_ind] num_parts = len(all_parts_mass.v) stars = Particles(num_parts) stars.mass = all_parts_mass.v | u.MSun stars.position = all_parts_pos.v | u.cm stars.velocity = all_parts_vel.v | u.cm/u.s stars.index_of_the_particle = range(num_parts) return stars
def get_ps_from_hdf5(partfile, particle_type='all'): """ Returns an AMUSE particle set from a FLASH hdf5 particle file (NOT a plot file) or checkpoint file. Keyword arguments: partfile -- the FLASH particle file particle_type -- the type of particle that you want from FLASH. Can be: star -- active FLASH particle type sink -- sink FLASH particle type any other string or none returns all particle types. The default is all. Returns: stars -- An AMUSE particle set. """ import h5py ds = h5py.File(partfile, 'r') part_data = ds.get('tracer particles') part_names = ds.get('particle names') part_data = np.array(part_data) part_names = np.array(part_names) real_scalars = ds.get('real scalars') part_time = real_scalars[0][1] if (part_data.size > 1): mass_ind = np.where(np.char.strip(part_names)=='mass')[0] tag_ind = np.where(np.char.strip(part_names)=='tag')[0] type_ind = np.where(np.char.strip(part_names)=='type')[0] ct_ind = np.where(np.char.strip(part_names)=='creation_time')[0] posx_ind = np.where(np.char.strip(part_names)=='posx')[0] posy_ind = np.where(np.char.strip(part_names)=='posy')[0] posz_ind = np.where(np.char.strip(part_names)=='posz')[0] velx_ind = np.where(np.char.strip(part_names)=='velx')[0] vely_ind = np.where(np.char.strip(part_names)=='vely')[0] velz_ind = np.where(np.char.strip(part_names)=='velz')[0] # Get the particle types type_part = part_data[:,type_ind].flatten() if (particle_type == 'star' or particle_type == 'stars'): # Find only star particles (particle_type=1) star_ind = np.where(type_part==1)[0] elif (particle_type == 'sink' or particle_type == 'sinks'): # Find only sink particles (particle_type=2) star_ind = np.where(type_part==2)[0] else: star_ind = np.arange(len(type_part)) num_parts = len(star_ind) # allocate the array stars = Particles(num_parts) # Masses of stars stars.mass = part_data[star_ind,mass_ind] | u.g # Tags of stars stars.tag = part_data[star_ind,tag_ind] # Creation times of stars stars.ct = part_data[star_ind,ct_ind] | u.s # Positions stars.x = part_data[star_ind,posx_ind] | u.cm stars.y = part_data[star_ind,posy_ind] | u.cm stars.z = part_data[star_ind,posz_ind] | u.cm # Velocities stars.vx = part_data[star_ind,velx_ind] | u.cm/u.s stars.vy = part_data[star_ind,vely_ind] | u.cm/u.s stars.vz = part_data[star_ind,velz_ind] | u.cm/u.s # Set an index attribute. Useful for some functions/codes in AMUSE. stars.index_of_the_particle = np.arange(num_parts) else: print "Error: No particles found in this file!" stars = None return stars