def run(h5_file= None, step=2.5, gpu_rank=0):
sp = SupercellBuilder(cif_path, verbose=False, debug=False)
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,198.0]),
projec_1=y_dir, projec_2=z_dir)
en = 100 # keV
semi_angle= 4e-3 # radians
max_ang = 200e-3
msa = MSAMPI(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 = 3.135 #
msa.build_potential_slices(slice_thickness)
print('rank %d: time to build atomic potential: %2.2f' % (comm_rank, 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()
if h5_file is not None:
msa.multislice(h5_write=h5_file)
else:
msa.multislice()
print('rank %02d: time to propagate probes:%2.2f' % (comm_rank, time() - t))
if comm_rank == 0: print('Simulation Finished Successfully!')
return
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 main(cifdir_path, outdir_path):
t = time()
cifpaths = get_cif_paths(cifdir_path)
num_sims = cifpaths.size
#num_sims = comm_size * 2
# HDF5
h5path = os.path.join(outdir_path, 'params_batch_%d.h5' % comm_rank)
mode = 'w'
with h5py.File(h5path, mode=mode) as f:
for (idx,
cif_path) in enumerate(cifpaths[comm_rank:num_sims:comm_size]):
manual = idx < (num_sims - comm_size)
spgroup_num, matname = parse_cif_path(cif_path)
sp = SupercellBuilder(cif_path, verbose=False, debug=False)
if comm_rank == 0 and bool(idx % 100):
print('time=%3.2f, num_sims= %d' %
(time() - t, idx * comm_size))
def main_test(cifdir_path):
cifpath_list = get_cif_paths(cifdir_path)
idx = np.random.randint(0, len(cifpath_list))
for _ in range(1000):
cif_path = cifpath_list[idx]
spgroup_num, matname = parse_cif_path(cif_path)
sp = SupercellBuilder(cif_path, verbose=False, debug=False)
sim_params = set_sim_params(sp,
energy=100e3,
orientation_num=3,
beam_overlap=2)
y_dir, z_dir = sim_params['y_dirs'][0], sim_params['z_dirs'][0]
sp.build_unit_cell()
sp.make_orthogonal_supercell(supercell_size=np.array(
[2 * 34.6, 2 * 34.6, 198.0]),
projec_1=y_dir,
projec_2=z_dir)
print(
"rank=%d, spgroup= %s, material=%s, z_dir=%s, y_dir=%s, d_hkl=%2.2f, semi_angle=%2.2f"
% (comm_rank, spgroup_num, matname, z_dir, y_dir,
sim_params['d_hkl'][0], sim_params['semi_angles'][0]))
if comm_rank == 0:
print('current idx: %d' % idx)
def main():
cifpath = 'cif_files/mvc-9418.cif'
spbuild = SupercellBuilder(cifpath, verbose=True)
spbuild.transform_unit_cell()
spbuild.build_unit_cell()
spbuild.make_orthogonal_supercell(xyz='examples/test.xyz')
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, 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, 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 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