15077298886

import copy
import functools
import multiprocessing as mp
import sys
from pathlib import Path

import numpy as np
import numpy.typing as npt

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


class Vibrations:
    """TODO."""

    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])
        self._vibration_coords = []
        self._vibration_forces = []

    def get_instance_of_other_type(
        self, vibrations_type: str
    ) -> "AimsVibrations | VaspVibrations":
        if vibrations_type == "aims":
            new_vibration = AimsVibrations()
        elif vibrations_type == "vasp":
            new_vibration = VaspVibrations()
        else:
            msg = f"Unsupported vibration type: {vibrations_type}"
            raise ValueError(msg)

        new_vibration.__dict__ = self.__dict__

        return new_vibration

    @property
    def vibration_coors(self) -> list[npt.NDArray[np.float64]]:
        return self._vibration_coords

    @vibration_coors.setter
    def vibration_coors(self, vibration_coords: list[npt.NDArray[np.float64]]) -> None:
        self._vibration_coords = vibration_coords

    @property
    def vibration_forces(self) -> list[npt.NDArray[np.float64]]:
        return self._vibration_forces

    @vibration_forces.setter
    def vibration_forces(self, vibration_forces: list[npt.NDArray[np.float64]]) -> None:
        self._vibration_forces = vibration_forces

    @property
    def hessian(self) -> npt.NDArray[np.float64]:
        return self._hessian

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

    @property
    def eigenvalues(self) -> npt.NDArray[np.float64]:
        return self._eigenvalues

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

    @property
    def eigenvectors(self) -> npt.NDArray[np.float64]:
        return self._eigenvectors

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

    def get_displacements(self, displacement: float = 0.0025) -> list:
        """
        Apply 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.

        """  # noqa: D205
        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))

        return np.sqrt(mass_tensor * mass_tensor.T)

    def get_hessian(
        self, set_constrained_atoms_zero: bool = False
    ) -> npt.NDArray[np.float64]:
        """
        Calculate the Hessian from the forces.

        This includes the atomic 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  # noqa: N806
        H = np.zeros([N, N])  # noqa: N806

        if not np.allclose(self.coors, self.vibration_coors[0]):  # pyright:ignore
            raise ValueError(
                "The first entry in vibration_coords must be identical to the "
                "undispaced geometry."
            )

        coords_0 = self.vibration_coors[0].flatten()
        F_0 = self.vibration_forces[0].flatten()  # noqa: N806

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

        for c, F in zip(self.vibration_coors, self.vibration_forces):  # noqa: N806
            dF = F.flatten() - F_0  # noqa: N806
            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: npt.NDArray[np.float64] | None = None
    ) -> npt.NDArray[np.float64]:
        """
        Symmetrieses the Hessian by using the lower triangular matrix.

        Parameters
        ----------
        hessian : npt.NDArray[np.float64] | None, default=None
            Hessian matrix to be symmetrized. If None, the Hessian of the object

        Returns
        -------
        npt.NDArray[np.float64]
            Symmetrized Hessian matrix.
        """
        if hessian is None:
            hessian = copy.deepcopy(self.hessian)

        hessian_new = hessian + hessian.T

        all_inds = list(range(len(self) * 3))  # pyright:ignore

        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

    def get_eigenvalues_and_eigenvectors(
        self,
        hessian: npt.NDArray[np.float64] | None = None,
        only_real: bool = True,
        symmetrize_hessian: bool = True,
        eigenvectors_to_cartesian: bool = False,
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        """
        Get all eigenvalues and eigenvectors of the hessian.

        Note that the eigenvectors are mass weighted.

        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
            Mass weighted eigenvectors of the Hessian given as a list of numpy
            arrays, where each array is a normalized displacement for the
            corresponding eigenfrequency.
        """
        if symmetrize_hessian:
            hessian = self.get_symmetrized_hessian(hessian=hessian)
        elif hessian is None:
            hessian = copy.deepcopy(self.hessian)

        if not hasattr(self, "hessian") or hessian is None:
            raise ValueError("Hessian must be given to calculate the Eigenvalues!")

        M = 1 / self.get_mass_tensor()  # noqa: N806

        omega2, X = np.linalg.eig(M * hessian)  # noqa: N806

        # 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])  # noqa: N806

        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

        # Convert eigenvector to Cartesian coordinates
        if eigenvectors_to_cartesian:
            m = np.tile(np.sqrt(self.get_atomic_masses()), (3, 1)).T  # pyright:ignore

            for index in range(len(eigenvectors)):
                eigenvectors[index] /= m

        return omega2, eigenvectors

    def get_eigenvalues_in_Hz(  # noqa: N802
        self, omega2: npt.NDArray[np.float64] | None = None
    ) -> npt.NDArray[np.float64]:
        """
        Determine angular vibration frequencies in Hz.

        Parameters
        ----------
        omega2 : npt.NDArray[np.float64] | None, default=None
            Eigenvalues of the mass weighted hessian.

        Returns
        -------
        omega_SI : np.array
            Array of the eigenfrequencies in Hz.
        """
        if omega2 is None:
            omega2 = self.get_eigenvalues_and_eigenvectors()[0]

        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)
        )
        return omega * conversion

    def get_eigenvalues_in_inverse_cm(
        self, omega2: npt.NDArray[np.float64] | None = 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)  # noqa: N806
        return omega_SI * units.INVERSE_CM_IN_HZ / (2 * np.pi)

    def get_eigenvalues_in_eV(  # noqa: N802
        self, omega2: npt.NDArray[np.float64] | None = None
    ) -> npt.NDArray[np.float64]:
        omega_SI = self.get_eigenvalues_in_Hz(omega2=omega2)  # noqa: N806
        return omega_SI * units.PLANCK_CONSTANT / (2 * np.pi) / units.JOULE_IN_EV

    def get_atom_type_index(self) -> npt.NDArray[np.int64]:
        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]

        return np.array(atom_type_index, dtype=int)

    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")
            sys.exit()

        # Projection onto the wave vector
        for i in range(number_of_atoms):
            # Projection on atom
            if project_on_atom > -1 and 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,
        use_numba: bool = True,
    ) -> npt.NDArray:
        """
        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)

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

    def get_cross_spectrum(
        self,
        velocities: npt.NDArray[np.float64],
        index_pair: tuple,
        time_step: float,
        bootstrapping_blocks: int = 1,
        bootstrapping_overlap: int = 0,
        cutoff_at_last_maximum: bool = True,
        window_function: str = "hann",
    ) -> tuple[npt.NDArray, npt.NDArray]:
        """
        PLACEHOLDE.

        Parameters
        ----------
        velocities : npt.NDArray[np.float64]
            Velocity time series.
        index_pair : tuple
            Indices of the two vibration between which the corss spectrum
            should be calculated. For instance (1, 4).
        time_step : float
            Time step in the velocity time series.
        bootstrapping_blocks : int, optional
            DESCRIPTION. The default is 1.
        bootstrapping_overlap : int, optional
            DESCRIPTION. The default is 0.

        Returns
        -------
        frequencies : np.array
            DESCRIPTION.
        cross_spectrum : np.array
            DESCRIPTION.
        """
        index_0, index_1 = index_pair

        velocities_proj = self.get_normal_mode_decomposition(velocities)

        frequencies, cross_spectrum = vu.get_cross_spectrum(
            velocities_proj[:, index_0],
            velocities_proj[:, index_1],
            time_step,
            bootstrapping_blocks=bootstrapping_blocks,
            bootstrapping_overlap=bootstrapping_overlap,
            cutoff_at_last_maximum=cutoff_at_last_maximum,
            window_function=window_function,
        )

        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: int = 1,
        frequency_cutoff: float | None = None,
        dirname: Path = Path("cross_spectrum"),
    ) -> None:
        """TODO."""
        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=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 dirname.is_dir():
            dirname.mkdir()

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

        np.savetxt(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,  # pyright: ignore
            dirname=dirname,
        )

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


def _output_cross_spectrum(
    index: npt.NDArray,
    velocities_proj: npt.NDArray,
    time_step: float,
    use_mem: bool,
    bootstrapping_blocks: int,
    bootstrapping_overlap: int,
    model_order: int,
    cutoff: int,
    dirname: Path,
) -> 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=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(
        dirname / f"cross_spectrum_{index_0}_{index_1}.csv",
        cross_spectrum[:cutoff],
    )


class AimsVibrations(Vibrations, AimsGeometry):  # pyright: ignore
    """TODO."""

    def __init__(self, filename: str | None = None):
        Vibrations.__init__(self)
        AimsGeometry.__init__(self, filename=filename)


class VaspVibrations(Vibrations, VaspGeometry):  # pyright: ignore
    """TODO."""

    def __init__(self, filename: str | None = None):
        Vibrations.__init__(self)
        VaspGeometry.__init__(self, filename=filename)

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

maurergroup / dfttoolkit / 15077298886

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