Source code for atomate2.common.jobs.anharmonicity

"""Jobs for running anharmonicity quantification."""

from __future__ import annotations

import contextlib
import logging
from typing import TYPE_CHECKING

import numpy as np
from jobflow import Flow, Response, job
from phonopy import Phonopy
from pymatgen.core import Structure
from pymatgen.core.units import kb
from pymatgen.io.phonopy import get_phonopy_structure, get_pmg_structure
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer

from atomate2.aims.schemas.calculation import Calculation
from atomate2.aims.utils.units import omegaToTHz
from atomate2.common.schemas.anharmonicity import AnharmonicityDoc

if TYPE_CHECKING:
    from pathlib import Path
    from typing import Any

    from atomate2.aims.jobs.base import BaseAimsMaker
    from atomate2.common.schemas.phonons import ForceConstants, PhononBSDOSDoc
    from atomate2.forcefields.jobs import ForceFieldStaticMaker
    from atomate2.vasp.jobs.base import BaseVaspMaker


logger = logging.getLogger(__name__)


class ImaginaryModeError(Exception):
    """Exception raised when an imaginary mode is detected.

    Attributes
    ----------
    largest_mode:
        The largest eigenmode to check the sign of
    message:
        Explanation of the error
    """

    def __init__(self, largest_mode: float) -> None:
        self.largest_mode = largest_mode
        self.message = f"""Structure has imaginary modes: the largest optical
                        eigenmode {largest_mode} < 0.0001"""
        super().__init__(self.message)


[docs] @job def get_phonon_supercell( structure: Structure, supercell_matrix: np.ndarray, ) -> Structure: """Get the phonon supercell of a structure. Parameters ---------- structure: Structure The structure to get phonon supercell for supercell_matrix: np.ndarray Supercell matrix for the phonon Returns ------- Structure The phonopy structure """ cell = get_phonopy_structure(structure) phonon = Phonopy( cell, supercell_matrix, ) return get_pmg_structure(phonon.supercell)
[docs] def get_sigma_per_site( structure: Structure, forces_dft: np.ndarray, forces_harmonic: np.ndarray, ) -> list[tuple]: """Compute a site-resolved sigma^A. This will calculate a sigma^A value for every site in the structure. Parameters ---------- structure: Structure The structure to use for the calculation forces_dft: np.ndarray DFT calculated forces forces_harmonic: np.ndarray Harmonic approximation of the forces Returns ------- list[tuple] List of tuples in the form ({wyckoff symbol: sites}, sigma^A) for all the sites in the structure """ # Ensure that DFT and harmonic forces are in np format forces_dft = np.array(forces_dft) forces_harmonic = np.array(forces_harmonic) # Check shapes of forces if len(np.shape(forces_dft)) == 2: forces_dft = np.expand_dims(forces_dft, axis=0) forces_harmonic = np.expand_dims(forces_harmonic, axis=0) sg_analyzer = SpacegroupAnalyzer(structure) sym_dataset = sg_analyzer.get_symmetry_dataset() wycks = np.array(sym_dataset.wyckoffs) sites = np.array(structure.sites) if np.shape(forces_dft)[1] == 3 * structure.num_sites: sites = sites.repeat(3) wycks = wycks.repeat(3) unique_sites = np.unique(wycks) site_to_wyckoff: dict = {wyck: [] for wyck in unique_sites} for idx, val in enumerate(wycks): site_to_wyckoff[val].append(sites[idx]) sigma_sites = [] for s in unique_sites: mask = wycks == s f_dft = forces_dft[:, mask] f_ha = forces_harmonic[:, mask] f_anharm = f_dft - f_ha sigma_sites.append(calc_sigma_a(f_anharm, f_dft)) return [ ( { list(site_to_wyckoff.keys())[i]: site_to_wyckoff[ list(site_to_wyckoff.keys())[i] ] }, sigma_sites[i], ) for i in range(len(sigma_sites)) ]
[docs] def get_sigma_per_element( structure: Structure, forces_dft: np.ndarray, forces_harmonic: np.ndarray, ) -> list[tuple[str, float]]: """Compute an element-resolved sigma^A. Every element type in the structure will have a sigma^A value calculated for it. Parameters ---------- structure: Structure The structure to use for the calculation forces_dft: np.ndarray DFT calculated forces forces_harmonic: np.ndarray Harmonic approximation of the forces Returns ------- list[tuple[str, float]] List of tuples in the form (atom symbol, sigma^A) for all the atoms in the structure. """ # Ensure that DFT and harmonic forces are in np format forces_dft = np.array(forces_dft) forces_harmonic = np.array(forces_harmonic) # Check shapes of forces if len(np.shape(forces_dft)) == 2: forces_dft = np.expand_dims(forces_dft, axis=0) forces_harmonic = np.expand_dims(forces_harmonic, axis=0) atom_numbers = np.array([site.specie.number for site in structure.sites]) if np.shape(forces_dft)[1] == 3 * structure.num_sites: atom_numbers = atom_numbers.repeat(3) symbols = np.array([site.specie.name for site in structure.sites]) unique_atoms = np.unique(atom_numbers) sigma_atom = [] unique_symbols = [] for u in unique_atoms: # Find atoms of type u in the structure mask = atom_numbers == u # Add symbol for atom u to list of examined atoms unique_symbols.append(symbols[mask][0]) # Take forces belonging to this atom f_dft = forces_dft[:, mask] f_ha = forces_harmonic[:, mask] f_anharm = f_dft - f_ha # Calculate sigma^A for this atom sigma_atom.append(calc_sigma_a(f_anharm, f_dft)) return [(unique_symbols[i], sigma_atom[i]) for i in range(len(sigma_atom))]
[docs] def box_muller( eig_vals: np.ndarray, eig_vecs: np.ndarray, temp: float, rng: np.random.Generator, ) -> np.ndarray: """Get a normally distributed random variable displacement. Uses the Box-Muller transform (https://en.wikipedia.org/wiki/Box-Muller_transform) Parameters ---------- eig_vals: np.ndarray Vector of harmonic eigenvalues (with first 3 removed for translational modes) eig_vecs: np.ndarray Matrix of harmonic eigenvectors (with first 3 removed for translational modes) temp: float Temperature in K to find velocity and displacement at rng: np.random.Generator Seeded random number generator Returns ------- np.ndarray Array with iid normally distributed displacements """ n_eigvals = eig_vals.shape[0] spread = np.sqrt(-2.0 * np.log(1.0 - rng.random(size=n_eigvals))) # Assign amplitudes (A_s) and phases (phi_s) a_s = spread * (np.sqrt(temp * kb) / eig_vals) phi_s = 2.0 * np.pi * rng.random(size=n_eigvals) # Get displacement (not normalized by sqrt(masses) yet) return (a_s * np.cos(phi_s) * eig_vecs).sum(axis=2)
[docs] @job def displace_structure( phonon_supercell: Structure, force_constants: ForceConstants = None, temp: float = 300, one_shot: bool = True, seed: int | None = None, mode_resolved: bool = False, n_samples: int = 1, ) -> list[Structure]: """Calculate the displaced structure. Procedure defined in doi.org/10.1103/PhysRevB.94.075125. Displaced supercell = Original supercell + Displacements Parameters ---------- phonon_supercell: Structure Supercell to distort force_constants: ForceConstants Force constants calculated from phonopy temp: float Temperature (in K) to displace structure at one_shot: bool If false, uses a normally distributed random number for zeta. The default is true. seed: int | None Seed to use for the random number generator (only used if one_shot_approx == False) mode_resolved: bool If true, calculate the mode-resolved sigma^A. This is false by default. n_samples: int Number of samples to generate (must be >= 1). An error is raised if n_samples == 1 and mode_resolved == True Returns ------- list[Structure] List of displaced Pymatgen structures """ if n_samples != 1 and one_shot: raise ValueError( f"One-shot can't be used unless n_sample == 1 not {n_samples}." ) if n_samples < 1: raise ValueError(f"n_sample must be >= 1, not {n_samples}.") if n_samples == 1 and mode_resolved: raise ValueError("n_sample must be > 1 when mode_resolved == True.") rng = np.random.default_rng(seed) coords = phonon_supercell.cart_coords disp = np.zeros(coords.shape) masses = np.array([site.species.weight for site in phonon_supercell.sites]) dynamical_matrix = build_dynmat(force_constants, phonon_supercell) eig_val, eig_vec = get_eigens(dynamical_matrix) # Check for imaginary modes if eig_val[3] < 0.0001: raise ImaginaryModeError(eig_val[3]) eig_val = eig_val[3:] x_acs = eig_vec[:, 3:].reshape((-1, 3, len(eig_val))) # gauge eigenvectors: largest value always positive for ii in range(x_acs.shape[-1]): vec = x_acs[:, :, ii] max_arg = np.argmax(abs(vec)) x_acs[:, :, ii] *= np.sign(vec.flat[max_arg]) inv_sqrt_mass = masses ** (-0.5) if one_shot: zetas = (-1) ** np.arange(len(eig_val)) a_s = np.sqrt(temp * kb) / eig_val * zetas disp = (a_s * x_acs).sum(axis=2) * inv_sqrt_mass[:, None] return [ Structure( lattice=phonon_supercell.lattice, species=phonon_supercell.species, coords=coords + disp, coords_are_cartesian=True, ) ] samples = [] for _ in range(n_samples): disp = box_muller(eig_val, x_acs, temp, rng) * inv_sqrt_mass[:, None] samples.append( Structure( lattice=phonon_supercell.lattice, species=phonon_supercell.species, coords=coords + disp, coords_are_cartesian=True, ) ) return samples
[docs] def build_dynmat( force_constants: ForceConstants, structure: Structure, ) -> np.ndarray: """Build the dynamical matrix. Parameters ---------- force_constants: ForceConstants Force constants calculated by Phonopy structure: Structure Structure as a Pymatgen Structure object Returns ------- np.ndarray The dynamical matrix """ force_constants_2d = ( np.array(force_constants.force_constants) .swapaxes(1, 2) .reshape(2 * (len(structure) * 3,)) ) masses = np.array([site.species.weight for site in structure.sites]) rminv = (masses**-0.5).repeat(3) return force_constants_2d * rminv[:, None] * rminv[None, :]
[docs] def get_eigens(dynmat: np.ndarray) -> tuple[np.ndarray, np.ndarray]: """Calculate the eigenmodes and eigenfrequencies from the dynamical matrix. Parameters ---------- dynmat: np.ndarray Dynamical matrix Returns ------- tuple[np.ndarray, np.ndarray] Tuple with first element being the array of eigenmodes and the second being the matrix of eigenfrequencies """ eig_val, eig_vec = np.linalg.eigh(dynmat) eig_val = np.sqrt(eig_val) return (eig_val, eig_vec)
[docs] def get_sigma_a_per_mode( force_constants: ForceConstants, structure: Structure, dft_forces: np.ndarray, harmonic_forces: np.ndarray, ) -> list[tuple[float, float]]: """Calculate sigma^A for each mode. Parameters ---------- force_constants: ForceConstants Force constants calculated by Phonopy structure: Structure Structure as a Pymatgen Structure object dft_forces: np.ndarray Array of DFT forces harmonic_forces: np.ndarray Array of harmonically approximated forces Returns ------- list[tuple[float, float]] List of tuples in the form (mode frequency in THz, Sigma^A) for all the modes in the structure """ # Projecting the forces dynamical_matrix = build_dynmat(force_constants, structure) eig_val, eig_vec = get_eigens(dynamical_matrix) eig_val = eig_val[3:] * omegaToTHz # Projection matrix P p = eig_vec.T masses = np.array([site.species.weight for site in structure.sites]) m = masses ** (-0.5) # Project the forces dft_forces = m.reshape((len(structure), 1)) * dft_forces dft_proj = np.array([p @ (f.flatten()) for f in np.array(dft_forces)])[:, 3:] harmonic_forces = m.reshape((len(structure), 1)) * harmonic_forces harmonic_proj = np.array([p @ (f.flatten()) for f in np.array(harmonic_forces)])[ :, 3: ] # Calculate sigma^A for each mode mode_sigma_vals = [] for mode in np.unique(eig_val): dft_vals = np.array( [dft_proj[:, idx] for idx in range(len(eig_val)) if eig_val[idx] == mode] ).flatten() harmonic_vals = np.array( [ harmonic_proj[:, idx] for idx in range(len(eig_val)) if eig_val[idx] == mode ] ).flatten() anharmonic_vals = [ dft_force - harmonic_force for dft_force, harmonic_force in zip(dft_vals, harmonic_vals, strict=True) ] sigma_a = np.std(anharmonic_vals) / np.std(dft_vals) mode_sigma_vals.append((mode, sigma_a)) return mode_sigma_vals
[docs] @job def get_forces( force_constants: ForceConstants, phonon_supercell: Structure, displaced_structures: dict[str, list], ) -> list[np.ndarray]: """Calculate the DFT forces and harmonic forces. Parameters ---------- force_constants: ForceConstants The force constants calculated by phonopy phonon_supercell: Structure The supercell used for the phonon calculation displaced_structures: dict[str, list] The output of run_displacements Returns ------- list[np.ndarray] List of forces in the form [DFT forces, harmonic forces] """ force_constants_2d = np.swapaxes( force_constants.force_constants, 1, 2, ).reshape(2 * (len(phonon_supercell) * 3,)) if isinstance(displaced_structures["coords"][0], Calculation): displacements = [ disp_data.output.structure.cart_coords - phonon_supercell.cart_coords for disp_data in displaced_structures["coords"] ] else: displacements = [ np.array(disp_data) - phonon_supercell.cart_coords for disp_data in displaced_structures["coords"] ] harmonic_forces = [ (-force_constants_2d @ displacement.flatten()).reshape((-1, 3)) for displacement in displacements ] dft_forces = [np.array(disp_data) for disp_data in displaced_structures["forces"]] return [dft_forces, harmonic_forces]
[docs] @job def get_sigmas( dft_forces: np.ndarray, harmonic_forces: np.ndarray, force_constants: ForceConstants, structure: Structure, one_shot: bool, element_resolved: bool = False, mode_resolved: bool = False, site_resolved: bool = False, ) -> dict: """Get the desired sigma^A measures. Parameters ---------- dft_forces: np.ndarray Array of DFT forces harmonic_forces: np.ndarray Array of forces from harmonic model force_constants: ForceConstants ForceConstants object containing a list of force constants structure: Structure The structure to calculate sigma^A on one_shot: bool If true, find the one-shot approximation. If false, find the full sigma^A number. element_resolved: bool If true, resolve sigma^A to the different elements. Default is false. mode_resolved: bool If true, resolve sigma^A to the different modes. Default is false. site_resolved: bool If true, resolve sigma^A to the different sites. Default is false. Returns ------- dict[str, Union[float, list]] Dictionary in the form {sigma^A type: float/list with sigma^A values} that contains all the sigma^A values """ dft_forces = np.array(dft_forces) harmonic_forces = np.array(harmonic_forces) anharmonic_forces = np.array( [ dft_force - harmonic_force for dft_force, harmonic_force in zip( dft_forces, harmonic_forces, strict=True ) ] ) sigma_dict: dict[str, Any] = {} if mode_resolved: sigma_dict["mode-resolved"] = get_sigma_a_per_mode( force_constants, structure, dft_forces, harmonic_forces, ) if element_resolved: sigma_dict["element-resolved"] = get_sigma_per_element( structure, dft_forces, harmonic_forces, ) if site_resolved: sigma_dict["site-resolved"] = get_sigma_per_site( structure, dft_forces, harmonic_forces, ) if one_shot: sigma_dict["one-shot"] = calc_sigma_a(anharmonic_forces, dft_forces) else: sigma_dict["full"] = calc_sigma_a(anharmonic_forces, dft_forces) return sigma_dict
[docs] def calc_sigma_a( anharmonic_forces: np.ndarray, dft_forces: np.ndarray, ) -> float: """Calculate sigma^A. Parameters ---------- anharmonic_forces: np.ndarray Array of anharmonic forces (dft - harmonic) dft_forces: np.ndarray Array of DFT forces Returns ------- float sigma^A number """ return np.std(anharmonic_forces) / np.std(dft_forces)
[docs] @job(data=["forces", "displaced_structures"]) def run_displacements( displacements: list[Structure], phonon_supercell: Structure, force_eval_maker: BaseVaspMaker | ForceFieldStaticMaker | BaseAimsMaker = None, prev_dir: str | Path = None, prev_dir_argname: str = None, socket: bool = False, ) -> Flow: """Run displaced structures. Note, this job will replace itself with N displacement calculations, or a single socket calculation for all displacements. Parameters ---------- displacements: Sequence All displacements to calculate phonon_supercell: Structure The supercell used for the phonon calculation force_eval_maker : .BaseVaspMaker or .ForceFieldStaticMaker or .BaseAimsMaker A maker to use to generate dispacement calculations prev_dir: str or Path The previous working directory prev_dir_argname: str argument name for the prev_dir variable socket: bool If True use the socket-io interface to increase performance Returns ------- Flow Flow calculating DFT forces on displaced structures """ force_eval_jobs = [] outputs: dict[str, list] = { "coords": [], "forces": [], "uuids": [], "dirs": [], } force_eval_job_kwargs = {} if prev_dir is not None and prev_dir_argname is not None: force_eval_job_kwargs[prev_dir_argname] = prev_dir if socket: force_eval_job = force_eval_maker.make(displacements, **force_eval_job_kwargs) info = { "phonon_supercell": phonon_supercell, "displaced_structures": displacements, } force_eval_job.update_maker_kwargs( {"_set": {"write_additional_data->anharmonicity_info:json": info}}, dict_mod=True, ) force_eval_jobs.append(force_eval_job) outputs["coords"] = force_eval_job.output.calcs_reversed outputs["uuids"] = [force_eval_job.output.uuid] * len(displacements) outputs["dirs"] = [force_eval_job.output.dir_name] * len(displacements) outputs["forces"] = force_eval_job.output.output.all_forces else: for idx, displacement in enumerate(displacements): if prev_dir is not None: force_eval_job = force_eval_maker.make(displacement, prev_dir=prev_dir) else: force_eval_job = force_eval_maker.make(displacement) force_eval_job.append_name( f" anharmonicity quant. {idx + 1}/{len(displacements)}" ) # we will add some meta data info = { "phonon_supercell": phonon_supercell, "displaced_structure": displacement, } with contextlib.suppress(Exception): force_eval_job.update_maker_kwargs( {"_set": {"write_additional_data->anharmonicity_info:json": info}}, dict_mod=True, ) force_eval_jobs.append(force_eval_job) outputs["coords"].append(force_eval_job.output.structure.cart_coords) outputs["uuids"].append(force_eval_job.output.uuid) outputs["dirs"].append(force_eval_job.output.dir_name) outputs["forces"].append(force_eval_job.output.output.forces) displacement_flow = Flow(force_eval_jobs, outputs) return Response(replace=displacement_flow)
[docs] @job def store_results( sigma_dict: dict[str, list | float], phonon_doc: PhononBSDOSDoc, one_shot: bool, temp: float, n_samples: int, seed: int | None, ) -> AnharmonicityDoc: """ Store the results in a schema object. Parameters ---------- sigma_dict: dict[str, Union[list, float]] Dictionary of computed sigma^A values. Possible contents are full, one-shot, atom-resolved, and mode-resolved. phonon_doc: PhononBSDOSDoc Info from phonon workflow to be stored one_shot: bool Whether the one-shot approximation was used (true) or not (false) temp: float Temperature (in K) to displace structure at n_samples: int How many displaced structures to sample seed: int | None What random seed was used to displace structures Returns ------- AnharmonicityDoc Document containing information on workflow results """ return AnharmonicityDoc.from_phonon_doc_sigma( sigma_dict=sigma_dict, phonon_doc=phonon_doc, one_shot=one_shot, temp=temp, n_samples=n_samples, seed=seed, )