def simulate(filehandle, cif_path, idx=None, gpu_id=0, clean_up=False): # load cif and get sim params spgroup_num, matname = parse_cif_path(cif_path) sp = SupercellBuilder(cif_path, verbose=False, debug=False) sim_params = get_sim_params(sp, grid_steps=np.array([2, 2]), orientation_num=3) z_dirs = sim_params['z_dirs'] y_dirs = sim_params['y_dirs'] cell_dim = sim_params['cell_dim'] slab_t = sim_params['slab_t'] sim_params['space_group'] = spgroup_num sim_params['material'] = matname energies = [100, 125, 150, 175, 200] for (sample_idx, (y_dir, (z_idx, z_dir), energy)) in enumerate(product(y_dirs, enumerate(z_dirs), energies)): try: t = time() # build supercell sp.build_unit_cell() sp.make_orthogonal_supercell(supercell_size=np.array( [cell_dim, cell_dim, slab_t]), projec_1=y_dir, projec_2=z_dir) # set simulation params slice_thickness = sim_params['d_hkl'][z_idx] #energy = sim_params['energy'] semi_angle = sim_params['semi_angles'][z_idx] probe_params = sim_params['probe_params'] sampling = sim_params['sampling'] grid_steps = sim_params['grid_steps'] # simulate msa = MSAGPU(energy, semi_angle, sp.supercell_sites, sampling=sampling, verbose=False, debug=False) ctx = msa.setup_device(gpu_rank=gpu_id) msa.calc_atomic_potentials() msa.build_potential_slices(slice_thickness) msa.build_probe(probe_dict=probe_params) msa.generate_probe_positions(grid_steps=grid_steps) msa.plan_simulation() msa.multislice() # process cbed and potential mask = msa.bandwidth_limit_mask(sampling, radius=1. / 3).astype(np.bool) proj_potential = process_potential(msa.potential_slices, mask=mask, normalize=True, fp16=True) cbed = process_cbed(msa.probes, normalize=True, fp16=True) # update sim_params dict sim_params = update_sim_params(sim_params, msa_cls=msa, sp_cls=sp) # #del sp # # clean-up context and/or allocated memory if clean_up and ctx is not None: msa.clean_up(ctx=ctx, vars=msa.vars) del msa else: msa.clean_up(ctx=None, vars=msa.vars) del msa # check data integrity has_nan = np.all(np.isnan(cbed)) or np.all( np.isnan(proj_potential)) wrong_shape = cbed.shape != ( 4, 512, 512) or proj_potential.shape != (1, 512, 512) if has_nan or wrong_shape: print('rank=%d, skipped simulation=%s, index=%d, error=NaN' % (comm_rank, cif_path, sample_idx)) else: # write to h5 / tfrecords / lmdb g = filehandle.create_group('sample_%d_%d' % (idx, sample_idx)) g.attrs['space_group'] = np.string_(sim_params['space_group']) g.attrs['material'] = np.string_(sim_params['material']) g.attrs['cif_path'] = np.string_(cif_path) g.attrs['energy'] = energy dset = g.create_dataset('z_dir', data=z_dir) dset = g.create_dataset('y_dir', data=y_dir) g.attrs['time'] = time() - t g.attrs['semi_angle'] = sim_params['semi_angles'][z_idx] g.attrs['d_hkl'] = sim_params['d_hkl'][z_idx] filehandle.flush() print('rank=%d, finished simulation=%s, index=%d' % (comm_rank, cif_path, sample_idx)) except Exception as e: print("rank=%d, skipped simulation=%s, error=%s" % (comm_rank, cif_path, format(e))) finally: try: if clean_up and ctx is not None: msa.clean_up(ctx=ctx, vars=msa.vars) del msa else: msa.clean_up(ctx=None, vars=msa.vars) del msa except: pass
def simulate(filehandle, h5g, idx=None, gpu_id=0, record_names=["2d_potential_", "cbed_"], clean_up=False): try: # load cif and get sim params cif_path = h5g.attrs['cif_path'].decode('ascii') z_dir = [0, 0, 1] y_dir = np.array([[1, 0, 0], [0, 1, 0]])[np.random.randint(2)] #z_dir = h5g['z_dir'][()] #y_dir = h5g['y_dir'][()] slice_thickness = h5g.attrs['d_hkl'] #semi_angle = h5g.attrs['semi_angle'] semi_angle = 0.01 sampling = np.array([256, 256]) cell_dim = 50 slab_t = 200 energy = 100e3 grid_steps = np.array([8, 8]) probe_params = { 'smooth_apert': True, 'scherzer': False, 'apert_smooth': 30, 'aberration_dict': { 'C1': 0., 'C3': 0, 'C5': 0. }, 'spherical_phase': True } semi_angle = 0.01 energy = np.random.randint(60, 140) probe_params['aberration_dict']['C3'] = np.round( 10**(np.random.rand() * 7)) # build supercell sp = SupercellBuilder(cif_path, verbose=False, debug=False) slice_thickness = max( 1.0, min(5.0, get_slice_thickness(sp, direc=np.array([0, 0, 1])))) # filter out angles = np.array(sp.structure.lattice.angles) angles = np.round(angles).astype(np.int) cutoff = np.array([90, 90, 90]) tol = 3 cubic_cond = np.logical_not( np.logical_and(angles > cutoff - tol, angles < cutoff + tol)).any() hexag_cond_1 = np.logical_and(angles[:2] > cutoff[:2] - tol, angles[:2] < cutoff[:2] + tol).any() hexag_cond_2 = np.logical_and(angles[-1] > 120 - tol, angles[-1] < 120 + tol) hexag_cond = np.logical_not(hexag_cond_1 and hexag_cond_2) if cubic_cond: if hexag_cond: return False else: pass sp.build_unit_cell() sp.make_orthogonal_supercell(supercell_size=np.array( [cell_dim, cell_dim, slab_t]), projec_1=y_dir, projec_2=z_dir) # simulate msa = MSAGPU(energy, semi_angle, sp.supercell_sites, sampling=sampling, verbose=False, debug=False) ctx = msa.setup_device(gpu_rank=gpu_id) msa.calc_atomic_potentials() msa.build_potential_slices(slice_thickness) msa.build_probe(probe_dict=probe_params) msa.generate_probe_positions(grid_steps=grid_steps) msa.plan_simulation() msa.multislice(bandwidth=1.) # process cbed and potential #mask = msa.bandwidth_limit_mask(sampling, radius=1./3).astype(np.bool) proj_potential = process_potential(msa.potential_slices, sampling=sampling, mask=None, normalize=True, fp16=True) cbed = process_cbed(msa.probes, normalize=True, fp16=True) # del sp # clean-up context and/or allocated memory if clean_up and ctx is not None: msa.clean_up(ctx=ctx, vars=msa.vars) del msa else: msa.clean_up(ctx=None, vars=msa.vars) del msa # check data integrity has_nan = np.all(np.isnan(cbed)) or np.all(np.isnan(proj_potential)) true_cbed_shape = (np.prod(grid_steps), ) + tuple(sampling) true_pot_shape = (1, ) + tuple(sampling) wrong_shape = cbed.shape != true_cbed_shape or proj_potential.shape != true_pot_shape if has_nan or wrong_shape: print("rank=%d, skipped simulation=%s, error=NaN" % (comm_rank, cif_path)) return False else: # write to h5 / tfrecords / lmdb if isinstance(filehandle, lmdb.Transaction): write_lmdb(filehandle, idx, cbed, proj_potential, record_names=record_names) print('rank=%d, finished simulation=%s' % (comm_rank, cif_path)) return True except Exception as e: print("rank=%d, skipped simulation=%s, error=%s" % (comm_rank, cif_path, format(e))) return False finally: # clean-up context and/or allocated memory try: if clean_up and ctx is not None: msa.clean_up(ctx=ctx, vars=msa.vars) del msa else: msa.clean_up(ctx=None, vars=msa.vars) del msa except: pass
def simulate(filehandle, cif_path, gpu_id=0, clean_up=False): # load cif and get sim params spgroup_num, matname = parse_cif_path(cif_path) index = 0 sp = SupercellBuilder(cif_path, verbose=False, debug=False) sim_params = get_sim_params(sp) z_dir = sim_params['z_dirs'][index] y_dir = sim_params['y_dirs'][index] cell_dim = sim_params['cell_dim'] slab_t = sim_params['slab_t'] sim_params['space_group'] = spgroup_num sim_params['material'] = matname # build supercell sp.build_unit_cell() sp.make_orthogonal_supercell(supercell_size=np.array( [cell_dim, cell_dim, slab_t]), projec_1=y_dir, projec_2=z_dir) # set simulation params slice_thickness = sim_params['d_hkl'][index] energy = sim_params['energy'] semi_angle = sim_params['semi_angles'][index] probe_params = sim_params['probe_params'] sampling = sim_params['sampling'] grid_steps = sim_params['grid_steps'] # simulate msa = MSAGPU(energy, semi_angle, sp.supercell_sites, sampling=sampling, verbose=False, debug=False) ctx = msa.setup_device(gpu_rank=gpu_id) msa.calc_atomic_potentials() msa.build_potential_slices(slice_thickness) msa.build_probe(probe_dict=probe_params) msa.generate_probe_positions(grid_steps=grid_steps) msa.plan_simulation() msa.multislice() # process cbed and potential mask = msa.bandwidth_limit_mask(sampling, radius=1. / 3).astype(np.bool) proj_potential = process_potential(msa.potential_slices, mask=mask) # update sim_params dict sim_params = update_sim_params(sim_params, msa_cls=msa, sp_cls=sp) # write to h5 / tfrecords if isinstance(filehandle, h5py.Group): write_h5(filehandle, msa.probes, proj_potential, sim_params) else: write_tfrecord(filehandle, msa.probes, proj_potential, sim_params) print('rank=%d, simulation=%s' % (comm_rank, cif_path)) # clean-up context and/or allocated memory if clean_up and ctx is not None: msa.clean_up(ctx=ctx, vars=msa.vars) else: msa.clean_up(ctx=None, vars=msa.vars)
def simulate(filehandle, cif_path, idx=None, gpu_id=0, clean_up=False): # load cif and get sim params spgroup_num, matname = parse_cif_path(cif_path) index = 1 sp = SupercellBuilder(cif_path, verbose=False, debug=False) t = time() sim_params = get_sim_params(sp) #if comm_rank == 0: print('time to get params: %2.3f' % (time() - t)) z_dir = sim_params['z_dirs'][index] y_dir = sim_params['y_dirs'][index] cell_dim = sim_params['cell_dim'] slab_t = sim_params['slab_t'] sim_params['space_group'] = spgroup_num sim_params['material'] = matname # build supercell t = time() sp.build_unit_cell() sp.make_orthogonal_supercell(supercell_size=np.array( [cell_dim, cell_dim, slab_t]), projec_1=y_dir, projec_2=z_dir) print('time to build cell : %2.3f' % (time() - t)) # set simulation params slice_thickness = sim_params['d_hkl'][index] energy = sim_params['energy'] semi_angle = sim_params['semi_angles'][index] probe_params = sim_params['probe_params'] sampling = sim_params['sampling'] grid_steps = sim_params['grid_steps'] # simulate t = time() msa = MSAGPU(energy, semi_angle, sp.supercell_sites, sampling=sampling, verbose=False, debug=False) ctx = msa.setup_device(gpu_rank=gpu_id) msa.calc_atomic_potentials() msa.build_potential_slices(slice_thickness) msa.build_probe(probe_dict=probe_params) msa.generate_probe_positions(grid_steps=grid_steps) #if comm_rank == 0: print('time to create context + other sims inputs : %2.3f' % (time() - t)) t = time() msa.plan_simulation() msa.multislice() #if comm_rank == 0: print('time to simulate cbed : %2.3f' % (time() - t)) # process cbed and potential mask = msa.bandwidth_limit_mask(sampling, radius=1. / 3).astype(np.bool) proj_potential = process_potential(msa.potential_slices, mask=mask, normalize=True, fp16=True) cbed = process_cbed(msa.probes, normalize=True, fp16=True) # update sim_params dict sim_params = update_sim_params(sim_params, msa_cls=msa, sp_cls=sp) has_nan = np.all(np.isnan(cbed)) or np.all(np.isnan(proj_potential)) wrong_shape = cbed.shape != (1024, 512, 512) or proj_potential.shape != (1, 512, 512) if has_nan or wrong_shape: # clean-up context and/or allocated memory #print('rank=%d, found this many %d nan in cbed' %(comm_rank, np.where(np.isnan(cbed)==True)[0].size)) #print('rank=%d, found this many %d nan in proj_pot' %(comm_rank, np.where(np.isnan(proj_potential)==True)[0].size)) #print('rank=%d, found this many %d nan in raw cbed' %(comm_rank, np.where(np.isnan(msa.probes)==True)[0].size)) if clean_up and ctx is not None: msa.clean_up(ctx=ctx, vars=msa.vars) else: msa.clean_up(ctx=None, vars=msa.vars) return False else: # write to h5 / tfrecords / lmdb if isinstance(filehandle, h5py.Group): write_h5(filehandle, cbed, proj_potential, sim_params) elif isinstance(filehandle, lmdb.Transaction): #write_lmdb(filehandle, idx + index, cbed, proj_potential, sim_params) write_lmdb(filehandle, idx, cbed, proj_potential, sim_params) elif isinstance(filehandle, tf.python_io.TFRecordWriter): write_tfrecord(filehandle, cbed, proj_potential, sim_params) # clean-up context and/or allocated memory if clean_up and ctx is not None: msa.clean_up(ctx=ctx, vars=msa.vars) else: msa.clean_up(ctx=None, vars=msa.vars) return True
def simulate(filehandle, cif_path, idx= None, gpu_ctx=None, clean_up=False, record_names=["2d_potential_", "cbed_"]): # load cif and get sim params spgroup_num, matname = parse_cif_path(cif_path) index = 0 sp = SupercellBuilder(cif_path, verbose=False, debug=False) latts = np.array(sp.structure.lattice.abc) if np.any(latts >= 10.) or np.all(latts >= 7): print('rank=%d, skipped simulation=%s, latt. const. too large=%s' % (comm_rank, cif_path, format(latts))) return False, None angles = np.array(sp.structure.lattice.angles) angles = np.round(angles).astype(np.int) cutoff = np.array([90,90,90]) tol = 2 cubic_cond = np.logical_not(np.logical_and(angles > cutoff - tol, angles < cutoff + tol)).any() hexag_cond_1 = np.logical_and(angles[:2] > cutoff[:2] - tol, angles[:2] < cutoff[:2] + tol).any() hexag_cond_2 = np.logical_and(angles[-1] > 120 - tol, angles[-1] < 120 + tol) hexag_cond = np.logical_not(hexag_cond_1 and hexag_cond_2) if cubic_cond: if hexag_cond: print("rank=%d, skipped simulation=%s, msg=not cubic/hexagonal" % (comm_rank, cif_path.split('/')[-2:])) return False, None else: pass sim_params = get_sim_params(sp, slab_t= 100, cell_dim = 50, grid_steps=np.array([8,8]), orientation_num=5, sampling=np.array([256,256])) cell_dim = sim_params['cell_dim'] slab_t = sim_params['slab_t'] sim_params['space_group']= spgroup_num sim_params['material'] = matname energies = np.linspace(100,200,4) np.random.shuffle(energies) counter = 0 for energy in energies: for index in range(len(sim_params['z_dirs'])): # build supercell z_dir = sim_params['z_dirs'][index] y_dir = sim_params['y_dirs'][index] sp.build_unit_cell() sp.make_orthogonal_supercell(supercell_size=np.array([cell_dim,cell_dim,slab_t]), projec_1=y_dir, projec_2=z_dir) # set simulation params slice_thickness = sim_params['d_hkl'][index] # energy = sim_params['energy'] semi_angle= sim_params['semi_angles'][index] probe_params = sim_params['probe_params'] sampling = sim_params['sampling'] grid_steps = sim_params['grid_steps'] ####### mods ###### # slice_thickness = max(1.0, min(5.0, get_slice_thickness(sp, direc=np.array([0,0,1])))) # semi_angle = 0.01 # energy = np.random.randint(60,140) probe_params['aberration_dict']['C3'] = np.round(10**(np.random.rand()*7)) ################## # simulate msa = MSAGPU(energy, semi_angle, sp.supercell_sites, sampling=sampling, verbose=False, debug=False) msa.calc_atomic_potentials() msa.build_potential_slices(gpu_ctx, slice_thickness) msa.build_probe(probe_dict=probe_params) msa.generate_probe_positions(grid_steps=grid_steps) msa.plan_simulation() msa.multislice(bandwidth=2./3) # process cbed and potential mask = msa.bandwidth_limit_mask(sampling, radius=1./3).astype(np.bool) proj_potential = process_potential(msa.potential_slices, mask=mask, normalize=True, fp16=True) cbed = process_cbed(msa.probes, normalize=True, fp16=True, new_shape=(64,256,256)) # update sim_params dict sim_params = update_sim_params(sim_params, msa_cls=msa, sp_cls=sp) has_nan = np.all(np.isnan(cbed)) or np.all(np.isnan(proj_potential)) wrong_shape = cbed.shape != (64, 256, 256) or proj_potential.shape != (1, 128, 128) if has_nan or wrong_shape: # clean-up context and/or allocated memory print('rank=%d, found this many %d nan in cbed' %(comm_rank, np.where(np.isnan(cbed)==True)[0].size)) print('rank=%d, found this many %d nan in proj_pot' %(comm_rank, np.where(np.isnan(proj_potential)==True)[0].size)) pass else: # write to h5 / tfrecords / lmdb if isinstance(filehandle, h5py.Group): write_h5(filehandle, cbed, proj_potential, sim_params) elif isinstance(filehandle, lmdb.Transaction): write_lmdb(filehandle, idx + counter , cbed, proj_potential, record_names=record_names) print('rank=%d, wrote sim_index=%d' % (comm_rank, idx+counter)) counter += 1 # free-up gpu memory msa.clean_up(ctx=None, vars=msa.vars) return True, counter
def run(step=2.5, gpu_rank=0): sp = SupercellBuilder('XYZ_files/Si.cif') z_dir = np.array([1,1,1]) y_dir = np.array([1,0,0]) #sp.transform_unit_cell() sp.build_unit_cell() sp.make_orthogonal_supercell(supercell_size=np.array([4*34.6,4*34.6,489.0]), projec_1=y_dir, projec_2=z_dir) en = 100 # keV semi_angle= 4e-3 # radians max_ang = 200e-3 msa = MSAGPU(en, semi_angle, sp.supercell_sites, sampling=np.array([512,512])) #verbose=True, debug=False) msa.setup_device(gpu_rank=gpu_rank) t = time() msa.calc_atomic_potentials() slice_thickness = 2.0 # msa.build_potential_slices(slice_thickness) print('time to build atomic potential: %2.2f' % (time() - t)) aberration_params = {'C1':500., 'C3': 3.3e7, 'C5':44e7} probe_params = {'smooth_apert': True, 'scherzer': False, 'apert_smooth': 60, 'aberration_dict':aberration_params, 'spherical_phase': True} msa.build_probe(probe_dict=probe_params) t = time() msa.generate_probe_positions(probe_step=np.array([step,step]), probe_range=np.array([[0.25,0.75],[0.25,0.75]])) msa.plan_simulation() msa.multislice() print('time to propagate probes:%2.2f' % (time() - t))
def simulate(filehandle, cif_path, idx= None, gpu_ctx=None, clean_up=False): # load cif and get sim params spgroup_num, matname = parse_cif_path(cif_path) sp = SupercellBuilder(cif_path, verbose=False, debug=False) latts = np.array(sp.structure.lattice.abc) if np.any(latts >= 10.): print('rank=%d, skipped simulation=%s, latt. const. too large=%s' % (comm_rank, cif_path, format(latts))) return sim_params = get_sim_params(sp, grid_steps=np.array([8,8]), orientation_num=3) z_dirs = sim_params['z_dirs'] y_dirs = sim_params['y_dirs'] cell_dim = sim_params['cell_dim'] slab_t = sim_params['slab_t'] sim_params['space_group']= spgroup_num sim_params['material'] = matname energies = np.arange(100,200,10) # ctx = msa.setup_device(gpu_rank=gpu_id) for (sample_idx, energy) in enumerate(energies): try: cbed_stack = [] for y_dir, (z_idx, z_dir) in zip(y_dirs, enumerate(z_dirs)): t = time() # build supercell sp.build_unit_cell() sp.make_orthogonal_supercell(supercell_size=np.array([cell_dim,cell_dim,slab_t]), projec_1=y_dir, projec_2=z_dir) # set simulation params slice_thickness = sim_params['d_hkl'][z_idx] semi_angle= sim_params['semi_angles'][z_idx] probe_params = sim_params['probe_params'] sampling = sim_params['sampling'] grid_steps = sim_params['grid_steps'] # simulate msa = MSAGPU(energy, semi_angle, sp.supercell_sites, sampling=sampling, verbose=False, debug=False) # ctx = msa.setup_device(gpu_rank=gpu_id) msa.calc_atomic_potentials() msa.build_potential_slices(gpu_ctx, slice_thickness) msa.build_probe(probe_dict=probe_params) msa.generate_probe_positions(grid_steps=grid_steps) msa.plan_simulation() msa.multislice(bandwidth=1./2) # process cbed cbed = msa.probes.mean(0) pot_isnan = np.isnan(msa.potential_slices) # update sim_params dict sim_params = update_sim_params(sim_params, msa_cls=msa, sp_cls=sp) # # clean-up context and/or allocated memory msa.clean_up(ctx=None, vars=msa.vars) # if clean_up and ctx is not None: # msa.clean_up(ctx=ctx, vars=msa.vars) # # del msa # else: # msa.clean_up(ctx=None, vars=msa.vars) # # del msa # check data integrity # isnan = np.isnan(cbed) # print('cbed elements with nan', np.where(isnan == True)[0].size) has_nan = np.all(np.isnan(cbed)) # pot_isnan = np.isnan(msa.potential_slices) # print('pot elements with nan', np.where(pot_isnan == True)[0].size) # if has_nan: # # print(z_dirs,y_dirs,sim_params['semi_angles'],sim_params['abc'], energy) # print('') wrong_shape = cbed.shape != (512, 512) if has_nan: print('rank=%d, skipped simulation=%s, index=%d, error=NaN, abc=%s' % (comm_rank, cif_path, sample_idx, format(sim_params['abc']))) pass elif wrong_shape: print('rank=%d, skipped simulation=%s, index=%d, error=wrong cbed shape' % (comm_rank, cif_path, sample_idx)) pass else: cbed_stack.append(cbed) # write to h5 / tfrecords / lmdb if len(cbed_stack) != 3: pass else: cbed_stack = np.stack(cbed_stack) g = filehandle.create_group('sample_%d_%d' % (idx, sample_idx)) g.attrs['space_group'] = np.string_(sim_params['space_group']) g.attrs['material'] = np.string_(sim_params['material']) g.attrs['energy_keV'] = energy dset = g.create_dataset('cbed_stack', data=cbed_stack) g.attrs['z_dirs'] = np.dstack(z_dirs) g.attrs['y_dirs'] = np.dstack(y_dirs) g.attrs['semi_angles_rad'] = sim_params['semi_angles'] g.attrs['d_hkls_angstrom'] = sim_params['d_hkl'] g.attrs['abc_angstrom'] = sim_params['abc'] g.attrs['angles_degree'] = sim_params['angles'] g.attrs['formula'] = sim_params['formula'] filehandle.flush() print('rank=%d, finished simulation=%s, index=%d' % (comm_rank, cif_path, sample_idx)) except Exception as e: print("rank=%d, skipped simulation=%s, error=%s" % (comm_rank, cif_path, format(e))) finally: try: if clean_up and gpu_ctx is not None: msa.clean_up(ctx=gpu_ctx, vars=msa.vars) del msa else: msa.clean_up(ctx=None, vars=msa.vars) del msa except: pass