Source code for ase2sprkkr.asr.setup.displacements

"""Module for generating atomic structures with displaced atoms."""

from pathlib import Path
from asr.core import command, option, ASRResult


[docs] def get_displacement_folder(atomic_index, cartesian_index, displacement_sign, displacement): """Generate folder name from (ia, iv, sign, displacement).""" cartesian_symbol = 'xyz'[cartesian_index] displacement_symbol = ' +-'[displacement_sign] foldername = (f'{displacement}-{atomic_index}' f'-{displacement_symbol}{cartesian_symbol}') folder = Path('displacements') / foldername return folder
[docs] def create_displacements_folder(folder): folder.mkdir(parents=True, exist_ok=False)
[docs] def get_all_displacements(atoms): """Generate ia, iv, sign for all displacements.""" for ia in range(len(atoms)): for iv in range(3): for sign in [-1, 1]: yield (ia, iv, sign)
[docs] def displace_atom(atoms, ia, iv, sign, delta): new_atoms = atoms.copy() pos_av = new_atoms.get_positions() pos_av[ia, iv] += sign * delta new_atoms.set_positions(pos_av) return new_atoms
@command('asr.setup.displacements') @option('--displacement', help='How much to displace atoms.', type=float) @option('--copy-params', help='Copy params.json to displacement folders.', type=bool) def main(displacement: float = 0.01, copy_params: bool = True) -> ASRResult: """Generate atomic displacements. Generate atomic structures with displaced atoms. The generated atomic structures are written to 'structure.json' and put into a directory with the structure displacements/{displacement}-{atomic_index}-{displacement_symbol}{cartesian_symbol} Notice that all generated directories are a sub-directory of displacements/. """ from ase.io import read from ase.parallel import world structure = read('structure.json') folders = [] params = Path('params.json') if not params.is_file(): copy_params = False else: params_text = params.read_text() for ia, iv, sign in get_all_displacements(structure): folder = get_displacement_folder(ia, iv, sign, displacement) if world.rank == 0: create_displacements_folder(folder) new_structure = displace_atom(structure, ia, iv, sign, displacement) new_structure.write(folder / 'structure.json') folders.append(str(folder)) if copy_params and params.is_file() and world.rank == 0: (folder / 'params.json').write_text(params_text) world.barrier() return {'folders': folders}