Source code for atomate2.vasp.jobs.amset

"""Module defining jobs for combining AMSET and VASP calculations."""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np
from click.testing import CliRunner
from jobflow import Flow, Response, job
from pymatgen.core.tensors import symmetry_reduce
from pymatgen.transformations.standard_transformations import (
    DeformStructureTransformation,
)

from atomate2 import SETTINGS
from atomate2.common.files import get_zfile
from atomate2.utils.file_client import FileClient
from atomate2.utils.path import strip_hostname
from atomate2.vasp.jobs.base import BaseVaspMaker
from atomate2.vasp.jobs.core import HSEBSMaker, NonSCFMaker
from atomate2.vasp.sets.core import (
    HSEBSSetGenerator,
    HSEStaticSetGenerator,
    NonSCFSetGenerator,
    StaticSetGenerator,
)

if TYPE_CHECKING:
    from typing import Any

    from emmet.core.math import Vector3D
    from pymatgen.core import Structure

    from atomate2.vasp.sets.base import VaspInputGenerator


[docs] @dataclass class DenseUniformMaker(NonSCFMaker): """ Maker to perform a dense uniform non-self consistent field calculation. Parameters ---------- name : str The job name. input_set_generator : .VaspInputGenerator A generator used to make the input set. write_input_set_kwargs : dict Keyword arguments that will get passed to :obj:`.write_vasp_input_set`. copy_vasp_kwargs : dict Keyword arguments that will get passed to :obj:`.copy_vasp_outputs`. run_vasp_kwargs : dict Keyword arguments that will get passed to :obj:`.run_vasp`. task_document_kwargs : dict Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`. stop_children_kwargs : dict Keyword arguments that will get passed to :obj:`.should_stop_children`. write_additional_data : dict Additional data to write to the current directory. Given as a dict of {filename: data}. Note that if using FireWorks, dictionary keys cannot contain the "." character which is typically used to denote file extensions. To avoid this, use the ":" character, which will automatically be converted to ".". E.g. ``{"my_file:txt": "contents of the file"}``. """ name: str = "dense uniform" input_set_generator: VaspInputGenerator = field( default_factory=lambda: NonSCFSetGenerator( mode="uniform", reciprocal_density=1000, user_incar_settings={"LWAVE": True} ) )
[docs] @dataclass class StaticDeformationMaker(BaseVaspMaker): """ Maker to perform a static calculations on structural deformations. The main difference to a normal static calculation is that this will write an explicit KPOINTS file, rather than using KSPACING. This is because all deformations ultimately need to be on exactly the same k-point mesh dimensions Parameters ---------- name : str The job name. input_set_generator : .VaspInputGenerator A generator used to make the input set. write_input_set_kwargs : dict Keyword arguments that will get passed to :obj:`.write_vasp_input_set`. copy_vasp_kwargs : dict Keyword arguments that will get passed to :obj:`.copy_vasp_outputs`. run_vasp_kwargs : dict Keyword arguments that will get passed to :obj:`.run_vasp`. task_document_kwargs : dict Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`. stop_children_kwargs : dict Keyword arguments that will get passed to :obj:`.should_stop_children`. write_additional_data : dict Additional data to write to the current directory. Given as a dict of {filename: data}. Note that if using FireWorks, dictionary keys cannot contain the "." character which is typically used to denote file extensions. To avoid this, use the ":" character, which will automatically be converted to ".". E.g. ``{"my_file:txt": "contents of the file"}``. """ name: str = "static deformation" input_set_generator: VaspInputGenerator = field( default_factory=lambda: StaticSetGenerator( user_kpoints_settings={"reciprocal_density": 100}, user_incar_settings={"KSPACING": None}, ) )
[docs] @dataclass class HSEStaticDeformationMaker(BaseVaspMaker): """ Maker to perform a HSE06 static calculations on structural deformations. The main difference to a normal HSE06 static calculation is that this will write an explicit KPOINTS file, rather than using KSPACING. This is because all deformations ultimately need to be on exactly the same k-point mesh dimensions Parameters ---------- name : str The job name. input_set_generator : .VaspInputGenerator A generator used to make the input set. write_input_set_kwargs : dict Keyword arguments that will get passed to :obj:`.write_vasp_input_set`. copy_vasp_kwargs : dict Keyword arguments that will get passed to :obj:`.copy_vasp_outputs`. run_vasp_kwargs : dict Keyword arguments that will get passed to :obj:`.run_vasp`. task_document_kwargs : dict Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`. stop_children_kwargs : dict Keyword arguments that will get passed to :obj:`.should_stop_children`. write_additional_data : dict Additional data to write to the current directory. Given as a dict of {filename: data}. Note that if using FireWorks, dictionary keys cannot contain the "." character which is typically used to denote file extensions. To avoid this, use the ":" character, which will automatically be converted to ".". E.g. ``{"my_file:txt": "contents of the file"}``. """ name: str = "static deformation" input_set_generator: VaspInputGenerator = field( default_factory=lambda: HSEStaticSetGenerator( user_kpoints_settings={"reciprocal_density": 100}, user_incar_settings={"KSPACING": None}, ) )
[docs] @dataclass class HSEDenseUniformMaker(HSEBSMaker): """ Maker to perform a dense uniform non-self consistent field calculation. Parameters ---------- name : str The job name. input_set_generator : .VaspInputGenerator A generator used to make the input set. write_input_set_kwargs : dict Keyword arguments that will get passed to :obj:`.write_vasp_input_set`. copy_vasp_kwargs : dict Keyword arguments that will get passed to :obj:`.copy_vasp_outputs`. run_vasp_kwargs : dict Keyword arguments that will get passed to :obj:`.run_vasp`. task_document_kwargs : dict Keyword arguments that will get passed to :obj:`.TaskDoc.from_directory`. stop_children_kwargs : dict Keyword arguments that will get passed to :obj:`.should_stop_children`. write_additional_data : dict Additional data to write to the current directory. Given as a dict of {filename: data}. Note that if using FireWorks, dictionary keys cannot contain the "." character which is typically used to denote file extensions. To avoid this, use the ":" character, which will automatically be converted to ".". E.g. ``{"my_file:txt": "contents of the file"}``. """ name: str = "dense uniform" input_set_generator: VaspInputGenerator = field( default_factory=lambda: HSEBSSetGenerator( mode="uniform_dense", zero_weighted_reciprocal_density=1000, user_incar_settings={"LWAVE": True}, ) )
[docs] @job def run_amset_deformations( structure: Structure, symprec: float = SETTINGS.SYMPREC, prev_dir: str | Path | None = None, static_deformation_maker: BaseVaspMaker | None = None, ) -> Response: """ Run amset deformations. Note, this job will replace itself with N static calculations, where N is the number of deformations. Parameters ---------- structure : .Structure A pymatgen structure. symprec : float Symmetry precision used to reduce the number of deformations. Set to None for no symmetry reduction. prev_dir : str or Path or None A previous VASP directory to use for copying VASP outputs. static_deformation_maker : .BaseVaspMaker or None A VaspMaker to use to generate the static deformation jobs. Returns ------- List[str] The directory names of each deformation calculation. """ from amset.deformation.generation import get_deformations if static_deformation_maker is None: static_deformation_maker = StaticDeformationMaker() deformations = get_deformations(0.005) if symprec is not None: deformations = list(symmetry_reduce(deformations, structure, symprec=symprec)) statics = [] outputs = [] for idx, deformation in enumerate(deformations): # deform the structure dst = DeformStructureTransformation(deformation=deformation) deformed_structure = dst.apply_transformation(structure) # create the job static_job = static_deformation_maker.make( deformed_structure, prev_dir=prev_dir ) static_job.append_name(f" {idx + 1}/{len(deformations)}") statics.append(static_job) # extract the outputs we want (only the dir name) outputs.append(static_job.output.dir_name) static_flow = Flow(statics, outputs) return Response(replace=static_flow)
[docs] @job def calculate_deformation_potentials( bulk_dir: str, deformation_dirs: list[str], symprec: float = SETTINGS.SYMPREC, ibands: tuple[list[int], list[int]] = None, ) -> dict[str, str]: """ Generate the deformation.h5 (containing deformation potentials) using AMSET. Note, this script just calls ``amset deform read``. Parameters ---------- bulk_dir : str The folder containing the bulk calculation data. deformation_dirs : list of str A list of folders for each deformation. symprec : float The symmetry precision used to reduce the number of deformations. Set to None if no-symmetry reduction was applied. ibands : tuple of list of int Which bands to include in the deformation.h5 file. Given as a tuple of one or two lists (one for each spin channel). The bands indices are zero indexed. Returns ------- dict A dictionary with the keys: - "dir_name": containing the directory where the deformation.h5 file was generated. - "log": The output log from ``amset deform read``. """ from amset.tools.deformation import read from click.testing import CliRunner # convert arguments into their command line equivalents # note, amset expects the band indices to be 1 indexed, whereas we store them # as zero indexed symprec_str = "N" if symprec is None else str(symprec) # TODO: Handle host names properly bulk_dir = strip_hostname(bulk_dir) deformation_dirs = [strip_hostname(d) for d in deformation_dirs] args = [ bulk_dir, *deformation_dirs, f"--symprec={symprec_str}", ] if ibands is not None: bands_str = ".".join( ",".join([str(idx + 1) for idx in band_ids]) for band_ids in ibands ) args.append(f"--bands={bands_str}") runner = CliRunner() result = runner.invoke(read, args, catch_exceptions=False) # TODO: Store some information about the deformation potentials, e.g., values # at CBM and VBM? return {"dir_name": str(Path.cwd()), "log": result.output}
[docs] @job def calculate_polar_phonon_frequency( structure: Structure, frequencies: list[float], eigenvectors: list[Vector3D], born_effective_charges: list[Vector3D], ) -> dict[str, list[float]]: """ Calculate the polar phonon frequency using amset. Parameters ---------- structure : .Structure A pymatgen structure. frequencies : list of float The phonon mode frequencies in THz. eigenvectors : list of list of float The phonon eigenvectors. born_effective_charges : list of list of float The born effective charges. Returns ------- dict A dictionary with the keys: - "frequency" (float): The polar phonon frequency. - "frequencies" (list[float]): A list of all phonon frequencies. - "weights" (list[float]): A list of weights for the frequencies. """ from amset.tools.phonon_frequency import calculate_effective_phonon_frequency effective_frequency, weights = calculate_effective_phonon_frequency( np.array(frequencies), np.array(eigenvectors), np.array(born_effective_charges), structure, ) return { "frequency": effective_frequency, "weights": weights.tolist(), "frequencies": frequencies, }
[docs] @job def generate_wavefunction_coefficients(dir_name: str) -> dict[str, Any]: """ Generate wavefunction.h5 file using amset. Parameters ---------- dir_name : str Directory containing WAVECAR and vasprun.xml files (can be gzipped). Returns ------- dict A dictionary with the keys: - "dir_name" (str): containing the directory where the wavefunction.h5 file was generated. - "log" (str): The output log from ``amset wave``. - "ibands" (Tuple[List[int], ...]): The bands included in the wavefunction.h5 file. Given as a tuple of one or two lists (one for each spin channel). The bands indices are zero indexed. """ from amset.tools.wavefunction import wave dir_name = strip_hostname(dir_name) # TODO: Handle hostnames properly. fc = FileClient() files = fc.listdir(dir_name) vasprun_file = Path(dir_name) / get_zfile(files, "vasprun.xml") wavecar_file = Path(dir_name) / get_zfile(files, "WAVECAR") # wavecar can't be gzipped, so copy it to current directory and unzip it fc.copy(wavecar_file, wavecar_file.name) fc.gunzip(wavecar_file.name) args = ["--wavecar=WAVECAR", f"--vasprun={vasprun_file}"] runner = CliRunner() result = runner.invoke(wave, args, catch_exceptions=False) ibands = _extract_ibands(result.output) # remove WAVECAR from current directory fc.remove("WAVECAR") return {"dir_name": str(Path.cwd()), "log": result.output, "ibands": ibands}
def _extract_ibands(log: str) -> tuple[list[int], ...]: """ Extract ibands from an ``amset wave`` log. Note in the log the band indices are 1 indexed but this function returns zero indexed band indices. Parameters ---------- log : str The log from ``amset wave``. Returns ------- tuple of list of int The bands included in the wavefunction.h5 file. Given as a tuple of one or two lists (one for each spin channel). The bands indices are zero indexed. """ result_splits = log.split("\n") for i in range(len(result_splits)): if "Including bands" in result_splits[i]: # non-spin polarised result system min_band, max_band = result_splits[i].split()[-1].split("—") return (list(range(int(min_band) - 1, int(max_band))),) if "Including:" in result_splits[i]: amin_band, amax_band = result_splits[i + 1].split()[-1].split("—") bmin_band, bmax_band = result_splits[i + 2].split()[-1].split("—") aibands = list(range(int(amin_band) - 1, int(amax_band))) bibands = list(range(int(bmin_band) - 1, int(bmax_band))) if "up" in result_splits[i + 1]: # up listed first return aibands, bibands else: # noqa: RET505 # down listed first return bibands, aibands raise ValueError("Could not find ibands in log.")