Ejemplo n.º 1
0
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
Ejemplo n.º 2
0
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
Ejemplo n.º 3
0
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)
Ejemplo n.º 4
0
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
Ejemplo n.º 5
0
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
Ejemplo n.º 6
0
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)) 
Ejemplo n.º 7
0
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