def load_snapshot_cholla_distributed( n_snapshot, data_type, fields, inputDir, subgrid = None, show_progess = True, precision = np.float64): base_name = '.h5.' fileName = '{0}{1}0'.format( n_snapshot, base_name ) #Load The first file to get parameters fileName = inputDir + fileName file = h5.File( fileName, 'r' ) dims = file.attrs['dims'] nprocs = file.attrs['nprocs'] domain = file.attrs['domain'] header = {} h_keys = file.attrs.keys() for key in h_keys: data_key = file.attrs[key] if data_key.shape == (1,): header[key] = data_key[0] else: header[key] = data_key file.close() grid_size = dims proc_grid = nprocs box_size = domain domain = get_domain_block( proc_grid, box_size, grid_size ) if subgrid == None: subgrid = [ [0,grid_size[0]], [0,grid_size[1]], [0,grid_size[2]], ] data_snapshot = load_snapshot_data_distributed( n_snapshot, inputDir, data_type, fields, subgrid, domain, precision, proc_grid, show_progess=show_progess ) data_snapshot['header'] = header return data_snapshot
# Load Enzo File snapKey = '{0:03}'.format(nSnap) inFileName = 'DD0{0}/data0{0}'.format( snapKey) ds = yt.load( inDir + inFileName ) data = ds.all_data() h = ds.hubble_constant current_z = np.float(ds.current_redshift) current_a = 1./(current_z + 1) # Set Domain parameters Lbox = 50000. proc_grid = [ 4, 2, 2] box_size = [ Lbox, Lbox, Lbox ] grid_size = [ 512, 512, 512 ] # Get the domain decomposition domain = get_domain_block( proc_grid, box_size, grid_size ) # Generate Particles ICs fields_particles = ['mass', 'pos_x', 'pos_y', 'pos_z', 'vel_x', 'vel_y', 'vel_z' ] outputBaseName = '{0}_particles.h5'.format(nSnap) generate_ics_particles_distributed( fields_particles, domain, proc_grid, data, ds, outputDir, outputBaseName, current_a, current_z, h, get_pid_indices=True) # Get Hydro ICs data_grid = ds.covering_grid( level=0, left_edge=ds.domain_left_edge, dims=ds.domain_dimensions ) fields_hydro = [ 'density', 'momentum_x', 'momentum_y', 'momentum_z', 'GasEnergy', 'Energy'] #It has to be in this order outputBaseName = '{0}.h5'.format(nSnap) generate_ics_grid_distributed( fields_hydro, domain, proc_grid, data_grid, ds, outputDir, outputBaseName, current_a, current_z, h )
def expand_data_particles_to_cholla(proc_grid, box_size, grid_size, inFileName, outDir, outputBaseName): domain = get_domain_block(proc_grid, box_size, grid_size) print('\n Loading data file: ', inFileName) inFile = h5py.File(inFileName, 'r') current_a = inFile.attrs['current_a'] current_z = inFile.attrs['current_z'] data = inFile['dm'] pos_x = data['pos_z'][...] pos_y = data['pos_y'][...] pos_z = data['pos_x'][...] vel_x = data['vel_z'][...] vel_y = data['vel_y'][...] vel_z = data['vel_x'][...] mass = data['mass'][...] nPart = mass.shape[0] print(' Nparticles: ', nPart) inFile.close() nprocs = proc_grid[0] * proc_grid[1] * proc_grid[2] for pId in range(nprocs): outputFileName = outDir + outputBaseName + ".{0}".format(pId) print(' Writing h5 file: ', outputFileName) outFile = h5py.File(outputFileName, 'w') outFile.attrs['box_size'] = box_size outFile.attrs['current_a'] = current_a outFile.attrs['current_z'] = current_z xMin, xMax = domain[pId]['box']['x'] yMin, yMax = domain[pId]['box']['y'] zMin, zMax = domain[pId]['box']['z'] print(('{0} x[{1} , {2}] y[{3} , {4}] z[{5}, {6}]'.format( pId, xMin, xMax, yMin, yMax, zMin, aMax))) indx_x = np.where(((pos_x >= xMin) & (pos_x < xMax))) indx_y = np.where(((pos_y >= yMin) & (pos_y < yMax))) indx_z = np.where(((pos_z >= zMin) & (pos_z < zMax))) # indx = [ idx for idx in range(len(pos_x)) if ( ( idx in indx_x ) )] indx_x = set(np.where(((pos_x >= xMin) & (pos_x < xMax)))[0]) indx_y = set(np.where(((pos_y >= yMin) & (pos_y < yMax)))[0]) indx_z = set(np.where(((pos_z >= zMin) & (pos_z < zMax)))[0]) indx = indx_x.intersection(indx_y) indx = list(indx.intersection(indx_z)) n_local = len(indx) pos_x_l = pos_x[indx] pos_y_l = pos_y[indx] pos_z_l = pos_z[indx] vel_x_l = vel_x[indx] vel_y_l = vel_y[indx] vel_z_l = vel_z[indx] mass_l = mass[indx] print(' n_local: ', n_local) print(' Current_a: ', current_a) outFile.attrs['n_particles_local'] = n_local # outFile.attrs['N_DM_file'] = np.float(nPart) outFile.create_dataset('mass', data=mass_l) outFile.create_dataset('pos_x', data=pos_x_l) outFile.create_dataset('pos_y', data=pos_y_l) outFile.create_dataset('pos_z', data=pos_z_l) outFile.create_dataset('vel_x', data=vel_x_l * np.sqrt(current_a)) outFile.create_dataset('vel_y', data=vel_y_l * np.sqrt(current_a)) outFile.create_dataset('vel_z', data=vel_z_l * np.sqrt(current_a)) outFile.close() print('')