13360600347

import copy
import functools
import multiprocessing as mp
import os
from typing import List, Tuple, Union

import numpy as np
import numpy.typing as npt
from scipy.signal import argrelextrema

import dfttoolkit.utils.units as units
import dfttoolkit.utils.vibrations_utils as vu
from dfttoolkit.geometry import AimsGeometry, VaspGeometry


class Vibrations:
    def __init__(self):
        # TODO This is currently a placeholder and should be developed further
        self.wave_vector = np.array([0.0, 0.0, 0.0])

    def get_instance_of_other_type(self, vibrations_type):
        if vibrations_type == "aims":
            new_vibration = AimsVibrations()
        if vibrations_type == "vasp":
            new_vibration = VaspVibrations()

        new_vibration.__dict__ = self.__dict__
        return new_vibration

    @property
    def vibration_coords(self):
        return self._vibration_coords

    @vibration_coords.setter
    def vibration_coords(self, vibration_coords: List[npt.NDArray[np.float64]]):
        self._vibration_coords = vibration_coords

    @property
    def vibration_forces(self):
        return self._vibration_forces

    @vibration_forces.setter
    def vibration_forces(self, vibration_forces: List[npt.NDArray[np.float64]]):
        self._vibration_forces = vibration_forces

    @property
    def hessian(self):
        return self._hessian

    @hessian.setter
    def hessian(self, hessian: npt.NDArray[np.float64]):
        self._hessian = hessian

    @property
    def eigenvalues(self):
        return self._eigenvalues

    @eigenvalues.setter
    def eigenvalues(self, eigenvalues: npt.NDArray[np.float64]):
        self._eigenvalues = eigenvalues

    @property
    def eigenvectors(self):
        return self._eigenvectors

    @eigenvectors.setter
    def eigenvectors(self, eigenvectors: npt.NDArray[np.float64]):
        self._eigenvectors = eigenvectors

    def get_displacements(self, displacement: float = 0.0025) -> list:
        """
        Applies a given displacement for each degree of freedom of self and
        generates a new vibration with it.

        Parameters
        ----------
        displacement : float, default=0.0025
            Displacement for finte difference calculation of vibrations in
            Angstrom.

        Returns
        -------
        list
            List of geometries where atoms have been displaced.

        """
        geometries_displaced = [self]

        directions = [-1, 1]

        for i in range(self.n_atoms):  # pyright:ignore
            for dim in range(3):
                if self.constrain_relax[i, dim]:  # pyright:ignore
                    continue

                for direction in directions:
                    geometry_displaced = copy.deepcopy(self)
                    geometry_displaced.coords[i, dim] += (  # pyright:ignore
                        displacement * direction
                    )

                    geometries_displaced.append(geometry_displaced)

        return geometries_displaced

    def get_mass_tensor(self) -> npt.NDArray[np.float64]:
        """
        Determine a NxN tensor containing sqrt(m_i*m_j).

        Returns
        -------
        mass_tensor : np.array
            Mass tensor in atomic units.

        """
        mass_vector = [
            self.periodic_table.get_atomic_mass(s)  # pyright:ignore
            for s in self.species  # pyright:ignore
        ]
        mass_vector = np.repeat(mass_vector, 3)

        mass_tensor = np.tile(mass_vector, (len(mass_vector), 1))
        mass_tensor = np.sqrt(mass_tensor * mass_tensor.T)

        return mass_tensor

    def get_hessian(
        self, set_constrained_atoms_zero: bool = False
    ) -> npt.NDArray[np.float64]:
        """
        Calculates the Hessian from the forces. This includes the atmoic masses
        since F = m*a.

        Parameters
        ----------
        set_constrained_atoms_zero : bool, default=False
            Set elements in Hessian that code for constrained atoms to zero.

        Returns
        -------
        H : np.array
            Hessian.

        """
        N = len(self) * 3  # pyright:ignore
        H = np.zeros([N, N])

        assert np.allclose(
            self.coords, self.vibration_coords[0]  # pyright:ignore
        ), "The first entry in vibration_coords must be identical to the undispaced geometry."

        coords_0 = self.vibration_coords[0].flatten()
        F_0 = self.vibration_forces[0].flatten()

        n_forces = np.zeros(N, np.int64)

        for c, F in zip(self.vibration_coords, self.vibration_forces):
            dF = F.flatten() - F_0
            dx = c.flatten() - coords_0
            ind = np.argmax(np.abs(dx))
            n_forces[ind] += 1
            displacement = dx[ind]

            if np.abs(displacement) < 1e-5:
                continue

            H[ind, :] -= dF / displacement

        for row in range(H.shape[0]):
            if n_forces[row] > 0:
                H[row, :] /= n_forces[row]  # prevent div by zero for unknown forces

        if set_constrained_atoms_zero:
            constrained = self.constrain_relax.flatten()  # pyright:ignore
            H[constrained, :] = 0
            H[:, constrained] = 0

        self.hessian = H
        return H

    def get_symmetrized_hessian(self, hessian=None):
        """
        Symmetrieses the Hessian by using the lower triangular matrix

        Parameters
        ----------
        hessian : TYPE, default=None
            DESCRIPTION

        Returns
        -------
        None.

        """
        if hessian is None:
            hessian = copy.deepcopy(self.hessian)

        hessian_new = hessian + hessian.T

        all_inds = list(range(len(self) * 3))

        constrain = self.constrain_relax.flatten()  # pyright:ignore
        constrained_inds = [i for i, c in enumerate(constrain) if c]
        constrained_inds = np.array(constrained_inds)

        unconstrained_inds = np.array(list(set(all_inds) - set(constrained_inds)))

        for i in unconstrained_inds:
            for j in unconstrained_inds:
                hessian_new[i, j] *= 0.5

        return hessian_new
        # self.hessian = (h+np.transpose(h))/2

    def get_eigenvalues_and_eigenvectors(
        self,
        hessian: Union[npt.NDArray[np.float64], None] = None,
        only_real: bool = True,
        symmetrize_hessian: bool = True,
    ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        """
        This function is supposed to return all eigenvalues and eigenvectors of
        the matrix self.hessian

        Parameters
        ----------
        hessian : npt.NDArray[np.float64], optional
            Hessian. The default is None.
        only_real : bool, default=True
            Returns only real valued eigenfrequencies + eigenmodes
            (ATTENTION: if you want to also include instable modes, you have to
            symmetrize the hessian as provided below).
        symmetrize_hessian : bool, default=True
            Symmetrise the hessian only for this function (no global change).

        Returns
        -------
        omega2 : np.array
            Direct eigenvalues as squared angular frequencies instead of
            inverse wavelengths.
        eigenvectors : np.array
            List of numpy arrays, where each array is a normalized
            displacement for the corresponding eigenfrequency, such that
            new_coords = coords + displacement * amplitude..

        """

        if symmetrize_hessian:
            hessian = self.get_symmetrized_hessian(hessian=hessian)
        elif hessian is None:
            hessian = copy.deepcopy(self.hessian)

        assert (
            hasattr(self, "hessian") and hessian is not None
        ), "Hessian must be given to calculate the Eigenvalues!"

        M = 1 / self.get_mass_tensor()

        omega2, X = np.linalg.eig(M * hessian)

        # only real valued eigen modes
        if only_real:
            real_mask = np.isreal(omega2)
            min_omega2 = 1e-3
            min_mask = omega2 >= min_omega2
            mask = np.logical_and(real_mask, min_mask)

            omega2 = np.real(omega2[mask])
            X = np.real(X[:, mask])

        eigenvectors = [column.reshape(-1, 3) for column in X.T]

        # sort modes by energies (ascending)
        ind_sort = np.argsort(omega2)
        eigenvectors = np.array(eigenvectors)[ind_sort, :, :]
        omega2 = omega2[ind_sort]

        self.eigenvalues = omega2
        self.eigenvectors = eigenvectors

        return omega2, eigenvectors

    def get_eigenvalues_in_Hz(
        self, omega2: Union[None, npt.NDArray[np.float64]] = None
    ) -> npt.NDArray[np.float64]:
        """
        Determine vibration frequencies in cm^-1.

        Parameters
        ----------
        omega2 : Union[None, np.array]
            Eigenvalues of the mass weighted hessian.

        Returns
        -------
        omega_SI : np.array
            Array of the eigenfrequencies in Hz.

        """
        if omega2 is None:
            omega2 = self.eigenvalues

        omega = np.sign(omega2) * np.sqrt(np.abs(omega2))  # pyright:ignore

        conversion = np.sqrt(
            (units.EV_IN_JOULE) / (units.ATOMIC_MASS_IN_KG * units.ANGSTROM_IN_METER**2)
        )
        omega_SI = omega * conversion

        return omega_SI

    def get_eigenvalues_in_inverse_cm(
        self, omega2: Union[None, npt.NDArray[np.float64]] = None
    ) -> npt.NDArray[np.float64]:
        """
        Determine vibration frequencies in cm^-1.

        Parameters
        ----------
        omega2 : Union[None, np.array]
            Eigenvalues of the mass weighted hessian.

        Returns
        -------
        f_inv_cm : np.array
            Array of the eigenfrequencies in cm^(-1).

        """
        omega_SI = self.get_eigenvalues_in_Hz(omega2=omega2)
        f_inv_cm = omega_SI * units.INVERSE_CM_IN_HZ / (2 * np.pi)

        return f_inv_cm

    def get_atom_type_index(self):

        n_atoms = len(self)  # pyright:ignore

        # Tolerance for accepting equivalent atoms in super cell
        masses = self.get_mass_of_all_atoms()  # pyright:ignore
        tolerance = 0.001

        primitive_cell_inverse = np.linalg.inv(self.lattice_vectors)  # pyright:ignore

        atom_type_index = np.array([None] * n_atoms)
        counter = 0
        for i in range(n_atoms):
            if atom_type_index[i] is None:
                atom_type_index[i] = counter
                counter += 1
            for j in range(i + 1, n_atoms):
                coordinates_atom_i = self.coords[i]  # pyright:ignore
                coordinates_atom_j = self.coords[j]  # pyright:ignore

                difference_in_cell_coordinates = np.around(
                    (
                        np.dot(
                            primitive_cell_inverse.T,
                            (coordinates_atom_j - coordinates_atom_i),
                        )
                    )
                )

                projected_coordinates_atom_j = coordinates_atom_j - np.dot(
                    self.lattice_vectors.T,  # pyright:ignore
                    difference_in_cell_coordinates,
                )
                separation = pow(
                    np.linalg.norm(projected_coordinates_atom_j - coordinates_atom_i),
                    2,
                )

                if separation < tolerance and masses[i] == masses[j]:
                    atom_type_index[j] = atom_type_index[i]

        atom_type_index = np.array(atom_type_index, dtype=int)

        return atom_type_index

    def get_velocity_mass_average(
        self, velocities: npt.NDArray[np.float64]
    ) -> npt.NDArray[np.float64]:
        """
        Weighs velocities by atomic masses.

        Parameters
        ----------
        velocities : npt.NDArray[np.float64]

        Returns
        -------
        velocities_mass_average : np.array
            Velocities weighted by atomic masses.

        """
        velocities_mass_average = np.zeros_like(velocities)

        for i in range(velocities.shape[1]):
            velocities_mass_average[:, i, :] = velocities[:, i, :] * np.sqrt(
                self.get_atomic_masses()[i]  # pyright:ignore
            )

        return velocities_mass_average

    def project_onto_wave_vector(
        self,
        velocities: npt.NDArray[np.float64],
        wave_vector: npt.NDArray[np.float64],
        project_on_atom: int = -1,
    ) -> npt.NDArray[np.float64]:

        number_of_primitive_atoms = len(self)  # pyright:ignore
        number_of_atoms = velocities.shape[1]
        number_of_dimensions = velocities.shape[2]

        coordinates = self.coords  # pyright:ignore
        atom_type = self.get_atom_type_index()

        velocities_projected = np.zeros(
            (
                velocities.shape[0],
                number_of_primitive_atoms,
                number_of_dimensions,
            ),
            dtype=complex,
        )

        if wave_vector.shape[0] != coordinates.shape[1]:
            print("Warning!! Q-vector and coordinates dimension do not match")
            exit()

        # Projection onto the wave vector
        for i in range(number_of_atoms):
            # Projection on atom
            if project_on_atom > -1:
                if atom_type[i] != project_on_atom:
                    continue

            for k in range(number_of_dimensions):
                velocities_projected[:, atom_type[i], k] += velocities[
                    :, i, k
                ] * np.exp(-1j * np.dot(wave_vector, coordinates[i, :]))

        # Normalize the velocities
        number_of_primitive_cells = number_of_atoms / number_of_primitive_atoms
        velocities_projected /= np.sqrt(number_of_primitive_cells)

        return velocities_projected

    def get_normal_mode_decomposition(
        self,
        velocities: npt.NDArray[np.float64],
        use_numba: bool = True,
    ) -> npt.NDArray[np.float64]:
        """
        Calculate the normal-mode-decomposition of the velocities. This is done
        by projecting the atomic velocities onto the vibrational eigenvectors.
        See equation 10 in: https://doi.org/10.1016/j.cpc.2017.08.017

        Parameters
        ----------
        velocities : npt.NDArray[np.float64]
            Array containing the velocities from an MD trajectory structured in
            the following way:
            [number of time steps, number of atoms, number of dimensions].

        Returns
        -------
        velocities_projected : npt.NDArray[np.float64]
            Velocities projected onto the eigenvectors structured as follows:
            [number of time steps, number of frequencies]

        """
        velocities = np.array(velocities, dtype=np.complex128)

        velocities_mass_averaged = self.get_velocity_mass_average(velocities)

        # velocities_proj_0 = vibrations.project_onto_wave_vector(velocities, q_vector)

        velocities_projected = vu.get_normal_mode_decomposition(
            velocities_mass_averaged,
            self.eigenvectors,
            use_numba=use_numba,
        )

        return velocities_projected

    def get_cross_correlation_function(
        self,
        velocities: npt.NDArray[np.float64],
        time_step: float,
        bootstrapping_blocks: int = 1,
        bootstrapping_overlap: int = 0,
    ):
        """
        Parameters
        ----------
        velocities : npt.NDArray[np.float64]
            DESCRIPTION.
        time_step : float
            DESCRIPTION.
        bootstrapping_blocks : int, default=1
            DESCRIPTION
        bootstrapping_overlap : int, default=0
            DESCRIPTION

        Returns
        -------
        autocorr_t : np.array
            DESCRIPTION.
        autocorr : np.array
            DESCRIPTION.

        """
        velocities_proj = self.get_normal_mode_decomposition(velocities)

        n_points = len(self.eigenvectors)

        autocorr_t = None
        autocorr = {}

        for index_0 in range(n_points):
            for index_1 in range(n_points):
                autocorr_block = vu.get_cross_correlation_function(
                    velocities_proj[:, index_0],
                    velocities_proj[:, index_1],
                    bootstrapping_blocks=bootstrapping_blocks,
                    bootstrapping_overlap=bootstrapping_overlap,
                )

                autocorr_t_block = np.linspace(
                    0, len(autocorr_block) * time_step, len(autocorr_block)
                )

                autocorr_t = autocorr_t_block
                autocorr[(index_0, index_1)] = autocorr_block

        return autocorr_t, autocorr

    def get_cross_spectrum(
        self,
        velocities: npt.NDArray[np.float64],
        time_step: float,
        use_mem: bool = False,
        bootstrapping_blocks: int = 1,
        bootstrapping_overlap: int = 0,
        model_order: int = 15,
    ):

        velocities_proj = self.get_normal_mode_decomposition(velocities)

        n_points = len(self.eigenvectors)

        if use_mem:
            frequencies = vu.get_cross_spectrum_mem(
                velocities_proj[:, 0],
                velocities_proj[:, 0],
                time_step,
                model_order,
                n_freqs=int(len(velocities_proj)),
            )[0]
        else:
            frequencies = vu.get_cross_spectrum(
                velocities_proj[:, 0],
                velocities_proj[:, 0],
                time_step,
                bootstrapping_blocks=bootstrapping_blocks,
                bootstrapping_overlap=bootstrapping_overlap,
            )[0]

        cross_spectrum = np.zeros(
            (n_points, n_points, len(frequencies)), dtype=np.complex128
        )

        for index_0 in range(n_points):
            for index_1 in range(n_points):
                print(index_0, index_1)

                if use_mem:
                    cross_spectrum_ij = vu.get_cross_spectrum_mem(
                        velocities_proj[:, index_0],
                        velocities_proj[:, index_1],
                        time_step,
                        model_order,
                        n_freqs=int(len(velocities_proj)),
                    )[1]
                else:
                    cross_spectrum_ij = vu.get_cross_spectrum(
                        velocities_proj[:, index_0],
                        velocities_proj[:, index_1],
                        time_step,
                        bootstrapping_blocks=bootstrapping_blocks,
                        bootstrapping_overlap=bootstrapping_overlap,
                    )[1]

                cross_spectrum[index_0, index_1, :] = cross_spectrum_ij

        return frequencies, cross_spectrum

    def output_cross_spectrum(
        self,
        velocities: npt.NDArray[np.float64],
        time_step: float,
        use_mem: bool = False,
        bootstrapping_blocks: int = 1,
        bootstrapping_overlap: int = 0,
        model_order: int = 15,
        processes=1,
        frequency_cutoff=None,
        dirname="cross_spectrum",
    ):

        velocities_proj = self.get_normal_mode_decomposition(velocities)

        n_points = len(self.eigenvectors)

        if use_mem:
            frequencies = vu.get_cross_spectrum_mem(
                velocities_proj[:, 0],
                velocities_proj[:, 0],
                time_step,
                model_order,
                n_freqs=int(len(velocities_proj)),
            )[0]
        else:
            frequencies = vu.get_cross_spectrum(
                velocities_proj[:, 0],
                velocities_proj[:, 0],
                time_step,
                bootstrapping_blocks=bootstrapping_blocks,
                bootstrapping_overlap=bootstrapping_overlap,
            )[0]

        if not os.path.isdir(dirname):
            os.mkdir(dirname)

        cutoff = -1
        if frequency_cutoff is not None:
            f_inv_cm = frequencies * units.INVERSE_CM_IN_HZ
            L = f_inv_cm < frequency_cutoff
            cutoff = np.sum(L)

        np.savetxt(os.path.join(dirname, "frequencies.csv"), frequencies[:cutoff])

        index = []
        for index_0 in range(n_points):
            for index_1 in range(n_points):
                if index_0 < index_1:
                    continue

                index.append((index_0, index_1))

        func = functools.partial(
            _output_cross_spectrum,
            velocities_proj=velocities_proj,
            time_step=time_step,
            use_mem=use_mem,
            bootstrapping_blocks=bootstrapping_blocks,
            bootstrapping_overlap=bootstrapping_overlap,
            model_order=model_order,
            cutoff=cutoff,
            dirname=dirname,
        )

        with mp.Pool(processes) as pool:
            pool.map(func, index)

    def get_coupling_matrix(
        self,
        velocities: npt.NDArray[np.float64],
        time_step: float,
        bootstrapping_blocks: int = 1,
        bootstrapping_overlap: int = 0,
    ):

        omega = self.get_eigenvalues_in_Hz()

        frequencies, power_spectrum = self.get_cross_spectrum(
            velocities,
            time_step,
            bootstrapping_blocks=bootstrapping_blocks,
            bootstrapping_overlap=bootstrapping_overlap,
        )

        n_points = len(self.eigenvectors)
        coupling_matrix = np.zeros((n_points, n_points))

        for index_0 in range(n_points):
            for index_1 in range(n_points):

                power_spectrum_index = np.real(power_spectrum[index_0, index_1, :])

                max_f = argrelextrema(power_spectrum_index, np.greater)[0]

                index_search = np.max([index_0, index_1])

                f_couple = omega[index_search] / (2 * np.pi)

                coupling_index_0 = np.argmin(np.abs(frequencies[max_f] - f_couple))
                coupling_index = max_f[coupling_index_0]

                coupling_matrix[index_0, index_1] = power_spectrum_index[coupling_index]

        return coupling_matrix


def _output_cross_spectrum(
    index,
    velocities_proj,
    time_step,
    use_mem,
    bootstrapping_blocks,
    bootstrapping_overlap,
    model_order,
    cutoff,
    dirname,
) -> None:
    index_0 = index[0]
    index_1 = index[1]

    if use_mem:
        cross_spectrum = vu.get_cross_spectrum_mem(
            velocities_proj[:, index_0],
            velocities_proj[:, index_1],
            time_step,
            model_order,
            n_freqs=int(len(velocities_proj)),
        )[1]
    else:
        cross_spectrum = vu.get_cross_spectrum(
            velocities_proj[:, index_0],
            velocities_proj[:, index_1],
            time_step,
            bootstrapping_blocks=bootstrapping_blocks,
            bootstrapping_overlap=bootstrapping_overlap,
        )[1]

    np.savetxt(
        os.path.join(dirname, f"cross_spectrum_{index_0}_{index_1}.csv"),
        cross_spectrum[:cutoff],
    )


class AimsVibrations(Vibrations, AimsGeometry):
    def __init__(self, filename=None):
        Vibrations.__init__(self)
        AimsGeometry.__init__(self, filename=filename)


class VaspVibrations(Vibrations, VaspGeometry):
    def __init__(self, filename=None):
        Vibrations.__init__(self)
        VaspGeometry.__init__(self, filename=filename)

1	import copy	×
2	import functools	×
3	import multiprocessing as mp	×
4	import os	×
5	from typing import List, Tuple, Union	×
6
7	import numpy as np	×
8	import numpy.typing as npt	×
9	from scipy.signal import argrelextrema	×
10
11	import dfttoolkit.utils.units as units	×
12	import dfttoolkit.utils.vibrations_utils as vu	×
13	from dfttoolkit.geometry import AimsGeometry, VaspGeometry	×
14
15
16	class Vibrations:	×
17	def __init__(self):	×
18	# TODO This is currently a placeholder and should be developed further
19	self.wave_vector = np.array([0.0, 0.0, 0.0])	×
20
21	def get_instance_of_other_type(self, vibrations_type):	×
22	if vibrations_type == "aims":	×
23	new_vibration = AimsVibrations()	×
24	if vibrations_type == "vasp":	×
25	new_vibration = VaspVibrations()	×
26
27	new_vibration.__dict__ = self.__dict__	×
28	return new_vibration	×
29
30	@property	×
31	def vibration_coords(self):	×
32	return self._vibration_coords	×
33
34	@vibration_coords.setter	×
35	def vibration_coords(self, vibration_coords: List[npt.NDArray[np.float64]]):	×
36	self._vibration_coords = vibration_coords	×
37
38	@property	×
39	def vibration_forces(self):	×
40	return self._vibration_forces	×
41
42	@vibration_forces.setter	×
43	def vibration_forces(self, vibration_forces: List[npt.NDArray[np.float64]]):	×
44	self._vibration_forces = vibration_forces	×
45
46	@property	×
47	def hessian(self):	×
48	return self._hessian	×
49
50	@hessian.setter	×
51	def hessian(self, hessian: npt.NDArray[np.float64]):	×
52	self._hessian = hessian	×
53
54	@property	×
55	def eigenvalues(self):	×
56	return self._eigenvalues	×
57
58	@eigenvalues.setter	×
59	def eigenvalues(self, eigenvalues: npt.NDArray[np.float64]):	×
60	self._eigenvalues = eigenvalues	×
61
62	@property	×
63	def eigenvectors(self):	×
64	return self._eigenvectors	×
65
66	@eigenvectors.setter	×
67	def eigenvectors(self, eigenvectors: npt.NDArray[np.float64]):	×
68	self._eigenvectors = eigenvectors	×
69
70	def get_displacements(self, displacement: float = 0.0025) -> list:	×
71	"""
72	Applies a given displacement for each degree of freedom of self and
73	generates a new vibration with it.
74
75	Parameters
76	----------
77	displacement : float, default=0.0025
78	Displacement for finte difference calculation of vibrations in
79	Angstrom.
80
81	Returns
82	-------
83	list
84	List of geometries where atoms have been displaced.
85
86	"""
87	geometries_displaced = [self]	×
88
89	directions = [-1, 1]	×
90
91	for i in range(self.n_atoms): # pyright:ignore	×
92	for dim in range(3):	×
93	if self.constrain_relax[i, dim]: # pyright:ignore	×
94	continue	×
95
96	for direction in directions:	×
97	geometry_displaced = copy.deepcopy(self)	×
98	geometry_displaced.coords[i, dim] += ( # pyright:ignore	×
99	displacement * direction
100	)
101
102	geometries_displaced.append(geometry_displaced)	×
103
104	return geometries_displaced	×
105
106	def get_mass_tensor(self) -> npt.NDArray[np.float64]:	×
107	"""
108	Determine a NxN tensor containing sqrt(m_i*m_j).
109
110	Returns
111	-------
112	mass_tensor : np.array
113	Mass tensor in atomic units.
114
115	"""
116	mass_vector = [	×
117	self.periodic_table.get_atomic_mass(s) # pyright:ignore
118	for s in self.species # pyright:ignore
119	]
120	mass_vector = np.repeat(mass_vector, 3)	×
121
122	mass_tensor = np.tile(mass_vector, (len(mass_vector), 1))	×
123	mass_tensor = np.sqrt(mass_tensor * mass_tensor.T)	×
124
125	return mass_tensor	×
126
127	def get_hessian(	×
128	self, set_constrained_atoms_zero: bool = False
129	) -> npt.NDArray[np.float64]:
130	"""
131	Calculates the Hessian from the forces. This includes the atmoic masses
132	since F = m*a.
133
134	Parameters
135	----------
136	set_constrained_atoms_zero : bool, default=False
137	Set elements in Hessian that code for constrained atoms to zero.
138
139	Returns
140	-------
141	H : np.array
142	Hessian.
143
144	"""
145	N = len(self) * 3 # pyright:ignore	×
146	H = np.zeros([N, N])	×
147
148	assert np.allclose(	×
149	self.coords, self.vibration_coords[0] # pyright:ignore
150	), "The first entry in vibration_coords must be identical to the undispaced geometry."
151
152	coords_0 = self.vibration_coords[0].flatten()	×
153	F_0 = self.vibration_forces[0].flatten()	×
154
155	n_forces = np.zeros(N, np.int64)	×
156
157	for c, F in zip(self.vibration_coords, self.vibration_forces):	×
158	dF = F.flatten() - F_0	×
159	dx = c.flatten() - coords_0	×
160	ind = np.argmax(np.abs(dx))	×
161	n_forces[ind] += 1	×
162	displacement = dx[ind]	×
163
164	if np.abs(displacement) < 1e-5:	×
165	continue	×
166
167	H[ind, :] -= dF / displacement	×
168
169	for row in range(H.shape[0]):	×
170	if n_forces[row] > 0:	×
171	H[row, :] /= n_forces[row] # prevent div by zero for unknown forces	×
172
173	if set_constrained_atoms_zero:	×
174	constrained = self.constrain_relax.flatten() # pyright:ignore	×
175	H[constrained, :] = 0	×
176	H[:, constrained] = 0	×
177
178	self.hessian = H	×
179	return H	×
180
181	def get_symmetrized_hessian(self, hessian=None):	×
182	"""
183	Symmetrieses the Hessian by using the lower triangular matrix
184
185	Parameters
186	----------
187	hessian : TYPE, default=None
188	DESCRIPTION
189
190	Returns
191	-------
192	None.
193
194	"""
195	if hessian is None:	×
196	hessian = copy.deepcopy(self.hessian)	×
197
198	hessian_new = hessian + hessian.T	×
199
200	all_inds = list(range(len(self) * 3))	×
201
202	constrain = self.constrain_relax.flatten() # pyright:ignore	×
203	constrained_inds = [i for i, c in enumerate(constrain) if c]	×
204	constrained_inds = np.array(constrained_inds)	×
205
206	unconstrained_inds = np.array(list(set(all_inds) - set(constrained_inds)))	×
207
208	for i in unconstrained_inds:	×
209	for j in unconstrained_inds:	×
210	hessian_new[i, j] *= 0.5	×
211
212	return hessian_new	×
213	# self.hessian = (h+np.transpose(h))/2
214
215	def get_eigenvalues_and_eigenvectors(	×
216	self,
217	hessian: Union[npt.NDArray[np.float64], None] = None,
218	only_real: bool = True,
219	symmetrize_hessian: bool = True,
220	) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
221	"""
222	This function is supposed to return all eigenvalues and eigenvectors of
223	the matrix self.hessian
224
225	Parameters
226	----------
227	hessian : npt.NDArray[np.float64], optional
228	Hessian. The default is None.
229	only_real : bool, default=True
230	Returns only real valued eigenfrequencies + eigenmodes
231	(ATTENTION: if you want to also include instable modes, you have to
232	symmetrize the hessian as provided below).
233	symmetrize_hessian : bool, default=True
234	Symmetrise the hessian only for this function (no global change).
235
236	Returns
237	-------
238	omega2 : np.array
239	Direct eigenvalues as squared angular frequencies instead of
240	inverse wavelengths.
241	eigenvectors : np.array
242	List of numpy arrays, where each array is a normalized
243	displacement for the corresponding eigenfrequency, such that
244	new_coords = coords + displacement * amplitude..
245
246	"""
247
248	if symmetrize_hessian:	×
249	hessian = self.get_symmetrized_hessian(hessian=hessian)	×
250	elif hessian is None:	×
251	hessian = copy.deepcopy(self.hessian)	×
252
253	assert (	×
254	hasattr(self, "hessian") and hessian is not None
255	), "Hessian must be given to calculate the Eigenvalues!"
256
257	M = 1 / self.get_mass_tensor()	×
258
259	omega2, X = np.linalg.eig(M * hessian)	×
260
261	# only real valued eigen modes
262	if only_real:	×
263	real_mask = np.isreal(omega2)	×
264	min_omega2 = 1e-3	×
265	min_mask = omega2 >= min_omega2	×
266	mask = np.logical_and(real_mask, min_mask)	×
267
268	omega2 = np.real(omega2[mask])	×
269	X = np.real(X[:, mask])	×
270
271	eigenvectors = [column.reshape(-1, 3) for column in X.T]	×
272
273	# sort modes by energies (ascending)
274	ind_sort = np.argsort(omega2)	×
275	eigenvectors = np.array(eigenvectors)[ind_sort, :, :]	×
276	omega2 = omega2[ind_sort]	×
277
278	self.eigenvalues = omega2	×
279	self.eigenvectors = eigenvectors	×
280
281	return omega2, eigenvectors	×
282
283	def get_eigenvalues_in_Hz(	×
284	self, omega2: Union[None, npt.NDArray[np.float64]] = None
285	) -> npt.NDArray[np.float64]:
286	"""
287	Determine vibration frequencies in cm^-1.
288
289	Parameters
290	----------
291	omega2 : Union[None, np.array]
292	Eigenvalues of the mass weighted hessian.
293
294	Returns
295	-------
296	omega_SI : np.array
297	Array of the eigenfrequencies in Hz.
298
299	"""
300	if omega2 is None:	×
301	omega2 = self.eigenvalues	×
302
303	omega = np.sign(omega2) * np.sqrt(np.abs(omega2)) # pyright:ignore	×
304
305	conversion = np.sqrt(	×
306	(units.EV_IN_JOULE) / (units.ATOMIC_MASS_IN_KG * units.ANGSTROM_IN_METER**2)
307	)
308	omega_SI = omega * conversion	×
309
310	return omega_SI	×
311
312	def get_eigenvalues_in_inverse_cm(	×
313	self, omega2: Union[None, npt.NDArray[np.float64]] = None
314	) -> npt.NDArray[np.float64]:
315	"""
316	Determine vibration frequencies in cm^-1.
317
318	Parameters
319	----------
320	omega2 : Union[None, np.array]
321	Eigenvalues of the mass weighted hessian.
322
323	Returns
324	-------
325	f_inv_cm : np.array
326	Array of the eigenfrequencies in cm^(-1).
327
328	"""
329	omega_SI = self.get_eigenvalues_in_Hz(omega2=omega2)	×
330	f_inv_cm = omega_SI * units.INVERSE_CM_IN_HZ / (2 * np.pi)	×
331
332	return f_inv_cm	×
333
334	def get_atom_type_index(self):	×
335
336	n_atoms = len(self) # pyright:ignore	×
337
338	# Tolerance for accepting equivalent atoms in super cell
339	masses = self.get_mass_of_all_atoms() # pyright:ignore	×
340	tolerance = 0.001	×
341
342	primitive_cell_inverse = np.linalg.inv(self.lattice_vectors) # pyright:ignore	×
343
344	atom_type_index = np.array([None] * n_atoms)	×
345	counter = 0	×
346	for i in range(n_atoms):	×
347	if atom_type_index[i] is None:	×
348	atom_type_index[i] = counter	×
349	counter += 1	×
350	for j in range(i + 1, n_atoms):	×
351	coordinates_atom_i = self.coords[i] # pyright:ignore	×
352	coordinates_atom_j = self.coords[j] # pyright:ignore	×
353
354	difference_in_cell_coordinates = np.around(	×
355	(
356	np.dot(
357	primitive_cell_inverse.T,
358	(coordinates_atom_j - coordinates_atom_i),
359	)
360	)
361	)
362
363	projected_coordinates_atom_j = coordinates_atom_j - np.dot(	×
364	self.lattice_vectors.T, # pyright:ignore
365	difference_in_cell_coordinates,
366	)
367	separation = pow(	×
368	np.linalg.norm(projected_coordinates_atom_j - coordinates_atom_i),
369	2,
370	)
371
372	if separation < tolerance and masses[i] == masses[j]:	×
373	atom_type_index[j] = atom_type_index[i]	×
374
375	atom_type_index = np.array(atom_type_index, dtype=int)	×
376
377	return atom_type_index	×
378
379	def get_velocity_mass_average(	×
380	self, velocities: npt.NDArray[np.float64]
381	) -> npt.NDArray[np.float64]:
382	"""
383	Weighs velocities by atomic masses.
384
385	Parameters
386	----------
387	velocities : npt.NDArray[np.float64]
388
389	Returns
390	-------
391	velocities_mass_average : np.array
392	Velocities weighted by atomic masses.
393
394	"""
395	velocities_mass_average = np.zeros_like(velocities)	×
396
397	for i in range(velocities.shape[1]):	×
398	velocities_mass_average[:, i, :] = velocities[:, i, :] * np.sqrt(	×
399	self.get_atomic_masses()[i] # pyright:ignore
400	)
401
402	return velocities_mass_average	×
403
404	def project_onto_wave_vector(	×
405	self,
406	velocities: npt.NDArray[np.float64],
407	wave_vector: npt.NDArray[np.float64],
408	project_on_atom: int = -1,
409	) -> npt.NDArray[np.float64]:
410
411	number_of_primitive_atoms = len(self) # pyright:ignore	×
412	number_of_atoms = velocities.shape[1]	×
413	number_of_dimensions = velocities.shape[2]	×
414
415	coordinates = self.coords # pyright:ignore	×
416	atom_type = self.get_atom_type_index()	×
417
418	velocities_projected = np.zeros(	×
419	(
420	velocities.shape[0],
421	number_of_primitive_atoms,
422	number_of_dimensions,
423	),
424	dtype=complex,
425	)
426
427	if wave_vector.shape[0] != coordinates.shape[1]:	×
428	print("Warning!! Q-vector and coordinates dimension do not match")	×
429	exit()	×
430
431	# Projection onto the wave vector
432	for i in range(number_of_atoms):	×
433	# Projection on atom
434	if project_on_atom > -1:	×
435	if atom_type[i] != project_on_atom:	×
436	continue	×
437
438	for k in range(number_of_dimensions):	×
439	velocities_projected[:, atom_type[i], k] += velocities[	×
440	:, i, k
441	] * np.exp(-1j * np.dot(wave_vector, coordinates[i, :]))
442
443	# Normalize the velocities
444	number_of_primitive_cells = number_of_atoms / number_of_primitive_atoms	×
445	velocities_projected /= np.sqrt(number_of_primitive_cells)	×
446
447	return velocities_projected	×
448
449	def get_normal_mode_decomposition(	×
450	self,
451	velocities: npt.NDArray[np.float64],
452	use_numba: bool = True,
453	) -> npt.NDArray[np.float64]:
454	"""
455	Calculate the normal-mode-decomposition of the velocities. This is done
456	by projecting the atomic velocities onto the vibrational eigenvectors.
457	See equation 10 in: https://doi.org/10.1016/j.cpc.2017.08.017
458
459	Parameters
460	----------
461	velocities : npt.NDArray[np.float64]
462	Array containing the velocities from an MD trajectory structured in
463	the following way:
464	[number of time steps, number of atoms, number of dimensions].
465
466	Returns
467	-------
468	velocities_projected : npt.NDArray[np.float64]
469	Velocities projected onto the eigenvectors structured as follows:
470	[number of time steps, number of frequencies]
471
472	"""
473	velocities = np.array(velocities, dtype=np.complex128)	×
474
475	velocities_mass_averaged = self.get_velocity_mass_average(velocities)	×
476
477	# velocities_proj_0 = vibrations.project_onto_wave_vector(velocities, q_vector)
478
479	velocities_projected = vu.get_normal_mode_decomposition(	×
480	velocities_mass_averaged,
481	self.eigenvectors,
482	use_numba=use_numba,
483	)
484
485	return velocities_projected	×
486
487	def get_cross_correlation_function(	×
488	self,
489	velocities: npt.NDArray[np.float64],
490	time_step: float,
491	bootstrapping_blocks: int = 1,
492	bootstrapping_overlap: int = 0,
493	):
494	"""
495	Parameters
496	----------
497	velocities : npt.NDArray[np.float64]
498	DESCRIPTION.
499	time_step : float
500	DESCRIPTION.
501	bootstrapping_blocks : int, default=1
502	DESCRIPTION
503	bootstrapping_overlap : int, default=0
504	DESCRIPTION
505
506	Returns
507	-------
508	autocorr_t : np.array
509	DESCRIPTION.
510	autocorr : np.array
511	DESCRIPTION.
512
513	"""
514	velocities_proj = self.get_normal_mode_decomposition(velocities)	×
515
516	n_points = len(self.eigenvectors)	×
517
518	autocorr_t = None	×
519	autocorr = {}	×
520
521	for index_0 in range(n_points):	×
522	for index_1 in range(n_points):	×
523	autocorr_block = vu.get_cross_correlation_function(	×
524	velocities_proj[:, index_0],
525	velocities_proj[:, index_1],
526	bootstrapping_blocks=bootstrapping_blocks,
527	bootstrapping_overlap=bootstrapping_overlap,
528	)
529
530	autocorr_t_block = np.linspace(	×
531	0, len(autocorr_block) * time_step, len(autocorr_block)
532	)
533
534	autocorr_t = autocorr_t_block	×
535	autocorr[(index_0, index_1)] = autocorr_block	×
536
537	return autocorr_t, autocorr	×
538
539	def get_cross_spectrum(	×
540	self,
541	velocities: npt.NDArray[np.float64],
542	time_step: float,
543	use_mem: bool = False,
544	bootstrapping_blocks: int = 1,
545	bootstrapping_overlap: int = 0,
546	model_order: int = 15,
547	):
548
549	velocities_proj = self.get_normal_mode_decomposition(velocities)	×
550
551	n_points = len(self.eigenvectors)	×
552
553	if use_mem:	×
554	frequencies = vu.get_cross_spectrum_mem(	×
555	velocities_proj[:, 0],
556	velocities_proj[:, 0],
557	time_step,
558	model_order,
559	n_freqs=int(len(velocities_proj)),
560	)[0]
561	else:
562	frequencies = vu.get_cross_spectrum(	×
563	velocities_proj[:, 0],
564	velocities_proj[:, 0],
565	time_step,
566	bootstrapping_blocks=bootstrapping_blocks,
567	bootstrapping_overlap=bootstrapping_overlap,
568	)[0]
569
570	cross_spectrum = np.zeros(	×
571	(n_points, n_points, len(frequencies)), dtype=np.complex128
572	)
573
574	for index_0 in range(n_points):	×
575	for index_1 in range(n_points):	×
576	print(index_0, index_1)	×
577
578	if use_mem:	×
579	cross_spectrum_ij = vu.get_cross_spectrum_mem(	×
580	velocities_proj[:, index_0],
581	velocities_proj[:, index_1],
582	time_step,
583	model_order,
584	n_freqs=int(len(velocities_proj)),
585	)[1]
586	else:
587	cross_spectrum_ij = vu.get_cross_spectrum(	×
588	velocities_proj[:, index_0],
589	velocities_proj[:, index_1],
590	time_step,
591	bootstrapping_blocks=bootstrapping_blocks,
592	bootstrapping_overlap=bootstrapping_overlap,
593	)[1]
594
595	cross_spectrum[index_0, index_1, :] = cross_spectrum_ij	×
596
597	return frequencies, cross_spectrum	×
598
599	def output_cross_spectrum(	×
600	self,
601	velocities: npt.NDArray[np.float64],
602	time_step: float,
603	use_mem: bool = False,
604	bootstrapping_blocks: int = 1,
605	bootstrapping_overlap: int = 0,
606	model_order: int = 15,
607	processes=1,
608	frequency_cutoff=None,
609	dirname="cross_spectrum",
610	):
611
612	velocities_proj = self.get_normal_mode_decomposition(velocities)	×
613
614	n_points = len(self.eigenvectors)	×
615
616	if use_mem:	×
617	frequencies = vu.get_cross_spectrum_mem(	×
618	velocities_proj[:, 0],
619	velocities_proj[:, 0],
620	time_step,
621	model_order,
622	n_freqs=int(len(velocities_proj)),
623	)[0]
624	else:
625	frequencies = vu.get_cross_spectrum(	×
626	velocities_proj[:, 0],
627	velocities_proj[:, 0],
628	time_step,
629	bootstrapping_blocks=bootstrapping_blocks,
630	bootstrapping_overlap=bootstrapping_overlap,
631	)[0]
632
633	if not os.path.isdir(dirname):	×
634	os.mkdir(dirname)	×
635
636	cutoff = -1	×
637	if frequency_cutoff is not None:	×
638	f_inv_cm = frequencies * units.INVERSE_CM_IN_HZ	×
639	L = f_inv_cm < frequency_cutoff	×
640	cutoff = np.sum(L)	×
641
642	np.savetxt(os.path.join(dirname, "frequencies.csv"), frequencies[:cutoff])	×
643
644	index = []	×
645	for index_0 in range(n_points):	×
646	for index_1 in range(n_points):	×
647	if index_0 < index_1:	×
648	continue	×
649
650	index.append((index_0, index_1))	×
651
652	func = functools.partial(	×
653	_output_cross_spectrum,
654	velocities_proj=velocities_proj,
655	time_step=time_step,
656	use_mem=use_mem,
657	bootstrapping_blocks=bootstrapping_blocks,
658	bootstrapping_overlap=bootstrapping_overlap,
659	model_order=model_order,
660	cutoff=cutoff,
661	dirname=dirname,
662	)
663
664	with mp.Pool(processes) as pool:	×
665	pool.map(func, index)	×
666
667	def get_coupling_matrix(	×
668	self,
669	velocities: npt.NDArray[np.float64],
670	time_step: float,
671	bootstrapping_blocks: int = 1,
672	bootstrapping_overlap: int = 0,
673	):
674
675	omega = self.get_eigenvalues_in_Hz()	×
676
677	frequencies, power_spectrum = self.get_cross_spectrum(	×
678	velocities,
679	time_step,
680	bootstrapping_blocks=bootstrapping_blocks,
681	bootstrapping_overlap=bootstrapping_overlap,
682	)
683
684	n_points = len(self.eigenvectors)	×
685	coupling_matrix = np.zeros((n_points, n_points))	×
686
687	for index_0 in range(n_points):	×
688	for index_1 in range(n_points):	×
689
690	power_spectrum_index = np.real(power_spectrum[index_0, index_1, :])	×
691
692	max_f = argrelextrema(power_spectrum_index, np.greater)[0]	×
693
694	index_search = np.max([index_0, index_1])	×
695
696	f_couple = omega[index_search] / (2 * np.pi)	×
697
698	coupling_index_0 = np.argmin(np.abs(frequencies[max_f] - f_couple))	×
699	coupling_index = max_f[coupling_index_0]	×
700
701	coupling_matrix[index_0, index_1] = power_spectrum_index[coupling_index]	×
702
703	return coupling_matrix	×
704
705
706	def _output_cross_spectrum(	×
707	index,
708	velocities_proj,
709	time_step,
710	use_mem,
711	bootstrapping_blocks,
712	bootstrapping_overlap,
713	model_order,
714	cutoff,
715	dirname,
716	) -> None:
717	index_0 = index[0]	×
718	index_1 = index[1]	×
719
720	if use_mem:	×
721	cross_spectrum = vu.get_cross_spectrum_mem(	×
722	velocities_proj[:, index_0],
723	velocities_proj[:, index_1],
724	time_step,
725	model_order,
726	n_freqs=int(len(velocities_proj)),
727	)[1]
728	else:
729	cross_spectrum = vu.get_cross_spectrum(	×
730	velocities_proj[:, index_0],
731	velocities_proj[:, index_1],
732	time_step,
733	bootstrapping_blocks=bootstrapping_blocks,
734	bootstrapping_overlap=bootstrapping_overlap,
735	)[1]
736
737	np.savetxt(	×
738	os.path.join(dirname, f"cross_spectrum_{index_0}_{index_1}.csv"),
739	cross_spectrum[:cutoff],
740	)
741
742
743	class AimsVibrations(Vibrations, AimsGeometry):	×
744	def __init__(self, filename=None):	×
745	Vibrations.__init__(self)	×
746	AimsGeometry.__init__(self, filename=filename)	×
747
748
749	class VaspVibrations(Vibrations, VaspGeometry):	×
750	def __init__(self, filename=None):	×
751	Vibrations.__init__(self)	×
752	VaspGeometry.__init__(self, filename=filename)	×

maurergroup / dfttoolkit / 13360600347

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous