15065810745

import copy

import numpy as np
import numpy.typing as npt
from ase import units
from ase.io.trajectory import Trajectory

import dfttoolkit.utils.vibrations_utils as vu
from dfttoolkit.geometry import Geometry


class MDTrajectory:
    def __init__(self, filename: str, cutoff_end: int = 0):
        traj_0 = Trajectory(filename)
        self.traj = []
        for ind in range(len(traj_0) - cutoff_end):
            self.traj.append(traj_0[ind])

        self.geom = Geometry()
        self.geom.get_from_ase_atoms_object(self.traj[0])

    def __add__(self, other_traj):
        traj = copy.deepcopy(self)
        traj += other_traj
        return traj

        self.traj = self.traj + other_traj.traj
        return None

    def __iadd__(self, other_traj):
        self.traj = self.traj + other_traj.traj
        return self

    def __len__(self):
        return len(self.traj)

    def get_velocities(
        self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
    ) -> npt.NDArray[np.float64]:
        """
        Get atomic velocities along an MD trajectory.

        Parameters
        ----------
        steps : int, optional
            Read every nth step. The default is 1 -> all steps are read. If for
            instance steps=5 every 5th step is read.
        cutoff_start : int, optional
            Cutoff n stept at the beginning of the trajectory. The default is 0.
        cutoff_end : int, optional
            Cutoff n stept at the end of the trajectory. The default is 0.

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

        """
        velocities = []

        for ind in range(cutoff_start, len(self.traj) - cutoff_end):
            if ind % steps == 0.0:
                velocities_new = self.traj[ind].get_velocities()
                velocities.append(velocities_new)

        return np.array(velocities, dtype=np.float64)

    def get_velocities_mass_weighted(
        self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
    ) -> npt.NDArray[np.float64]:
        """
        Weighs velocities by atomic masses.

        Parameters
        ----------
        steps : int, optional
            Read every nth step. The default is 1 -> all steps are read. If for
            instance steps=5 every 5th step is read.
        cutoff_start : int, optional
            Cutoff n stept at the beginning of the trajectory. The default is 0.
        cutoff_end : int, optional
            Cutoff n stept at the end of the trajectory. The default is 0.

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

        """
        velocities = self.get_velocities(
            steps=steps, cutoff_start=cutoff_start, cutoff_end=cutoff_end
        )

        velocities_mass_weighted = np.zeros_like(velocities)
        atomic_masses = self.geom.get_atomic_masses()

        for i in range(velocities.shape[1]):
            velocities_mass_weighted[:, i, :] = velocities[:, i, :] * np.sqrt(
                atomic_masses[i]
            )

        return velocities_mass_weighted

    def get_velocities_projected(
        self,
        projection_vectors: npt.NDArray[np.float64],
        mass_weighted: bool = True,
        steps: int = 1,
        cutoff_start: int = 0,
        cutoff_end: int = 0,
        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
        ----------
        projection_vectors : npt.NDArray[np.float64]
            Array containing the vectors onto which the velocities should be
            projected. Normally you will want this to be the eigenvectors
            return by dfttoolkit.vibration.get_eigenvalues_and_eigenvectors.
        mass_weighted : bool
            Wheter the velocities should be mass weighted. If you do
            normal-mode-decomposition with eigenvectors than this should be
            True. The default is True.
        steps : int, optional
            Read every nth step. The default is 1 -> all steps are read. If for
            instance steps=5 every 5th step is read.
        cutoff_start : int, optional
            Cutoff n stept at the beginning of the trajectory. The default is 0.
        cutoff_end : int, optional
            Cutoff n stept at the end of the trajectory. The default is 0.
        use_numba : bool
            Use numba for projection. Usually faster, but if your numbe
            installation is wonky, it may produce falty results.

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

        """
        if mass_weighted:
            velocities = self.get_velocities_mass_weighted(
                steps=steps, cutoff_start=cutoff_start, cutoff_end=cutoff_end
            )
        else:
            velocities = self.get_velocities(
                steps=steps, cutoff_start=cutoff_start, cutoff_end=cutoff_end
            )

        return vu.get_normal_mode_decomposition(
            velocities,
            projection_vectors,
            use_numba=use_numba,
        )


    def get_kinetic_energies(
        self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
    ) -> npt.NDArray[np.float64]:
        """
        Weighs velocities by atomic masses.

        Parameters
        ----------
        steps : int, optional
            Read every nth step. The default is 1 -> all steps are read. If for
            instance steps=5 every 5th step is read.
        cutoff_start : int, optional
            Cutoff n stept at the beginning of the trajectory. The default is 0.
        cutoff_end : int, optional
            Cutoff n stept at the end of the trajectory. The default is 0.

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

        """
        velocities = self.get_velocities_mass_weighted(
            steps=steps, cutoff_start=cutoff_start, cutoff_end=cutoff_end
        )
        velocities_norm = np.linalg.norm(velocities, axis=2)

        return 0.5 * velocities_norm**2


    def get_temperature(
        self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
    ) -> npt.NDArray[np.float64]:
        """
        Get atomic temperature along an MD trajectory.

        Parameters
        ----------
        steps : int, optional
            Read every nth step. The default is 1 -> all steps are read. If for
            instance steps=5 every 5th step is read.
        cutoff_start : int, optional
            Cutoff n stept at the beginning of the trajectory. The default is 0.
        cutoff_end : int, optional
            Cutoff n stept at the end of the trajectory. The default is 0.

        Returns
        -------
        temperature : npt.NDArray[np.float64]

        """
        temperature = []

        for ind in range(cutoff_start, len(self.traj) - cutoff_end):
            if ind % steps == 0.0:
                atoms = self.traj[ind]
                unconstrained_atoms = len(atoms) - len(atoms.constraints[0].index)

                ekin = atoms.get_kinetic_energy() / unconstrained_atoms
                temperature.append(ekin / (1.5 * units.kB))

        return np.array(temperature, dtype=np.float64)

    def get_total_energy(
        self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
    ) -> npt.NDArray[np.float64]:
        """
        Get atomic total energy along an MD trajectory.

        Parameters
        ----------
        steps : int, optional
            Read every nth step. The default is 1 -> all steps are read. If for
            instance steps=5 every 5th step is read.
        cutoff_start : int, optional
            Cutoff n stept at the beginning of the trajectory. The default is 0.
        cutoff_end : int, optional
            Cutoff n stept at the end of the trajectory. The default is 0.

        Returns
        -------
        total_energy : npt.NDArray[np.float64]

        """
        total_energy = []

        for ind in range(cutoff_start, len(self.traj) - cutoff_end):
            if ind % steps == 0.0:
                atoms = self.traj[ind]
                total_energy.append(atoms.get_total_energy())

        return np.array(total_energy, dtype=np.float64)

    def get_coords(self, atoms):
        unconstrained_atoms = atoms.constraints[0].index

        coords = []
        for ind, atom in enumerate(atoms):
            if ind not in unconstrained_atoms:
                coords.append(atom.position)

        return np.array(coords)

    def get_atomic_displacements(
        self,
        coords_0=None,
        steps: int = 1,
        cutoff_start: int = 0,
        cutoff_end: int = 0,
    ) -> npt.NDArray[np.float64]:
        """
        Get atomic atomic displacements with respect to the first time step
        along an MD trajectory.

        Parameters
        ----------
        steps : int, optional
            Read every nth step. The default is 1 -> all steps are read. If for
            instance steps=5 every 5th step is read.
        cutoff_start : int, optional
            Cutoff n stept at the beginning of the trajectory. The default is 0.
        cutoff_end : int, optional
            Cutoff n stept at the end of the trajectory. The default is 0.

        Returns
        -------
        atomic_displacements : npt.NDArray[np.float64]

        """
        atomic_displacements = []

        if coords_0 is None:
            coords_0 = self.get_coords(self.traj[cutoff_start])

        for ind in range(cutoff_start, len(self.traj) - cutoff_end):
            if ind % steps == 0.0:
                coords = self.get_coords(self.traj[ind])

                disp = np.linalg.norm(coords - coords_0, axis=1)
                atomic_displacements.append(disp)

        return np.array(atomic_displacements, dtype=np.float64)

1	import copy	×
2
3	import numpy as np	×
4	import numpy.typing as npt	×
5	from ase import units	×
6	from ase.io.trajectory import Trajectory	×
7
8	import dfttoolkit.utils.vibrations_utils as vu	×
9	from dfttoolkit.geometry import Geometry	×
10
11
12	class MDTrajectory:	×
13	def __init__(self, filename: str, cutoff_end: int = 0):	×
14	traj_0 = Trajectory(filename)	×
15	self.traj = []	×
16	for ind in range(len(traj_0) - cutoff_end):	×
17	self.traj.append(traj_0[ind])	×
18
19	self.geom = Geometry()	×
20	self.geom.get_from_ase_atoms_object(self.traj[0])	×
21
22	def __add__(self, other_traj):	×
23	traj = copy.deepcopy(self)	×
24	traj += other_traj	×
25	return traj	×
26
27	self.traj = self.traj + other_traj.traj
28	return None
29
30	def __iadd__(self, other_traj):	×
31	self.traj = self.traj + other_traj.traj	×
32	return self	×
33
34	def __len__(self):	×
35	return len(self.traj)	×
36
37	def get_velocities(	×
38	self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
39	) -> npt.NDArray[np.float64]:
40	"""
41	Get atomic velocities along an MD trajectory.
42
43	Parameters
44	----------
45	steps : int, optional
46	Read every nth step. The default is 1 -> all steps are read. If for
47	instance steps=5 every 5th step is read.
48	cutoff_start : int, optional
49	Cutoff n stept at the beginning of the trajectory. The default is 0.
50	cutoff_end : int, optional
51	Cutoff n stept at the end of the trajectory. The default is 0.
52
53	Returns
54	-------
55	velocities : npt.NDArray[np.float64]
56
57	"""
58	velocities = []	×
59
60	for ind in range(cutoff_start, len(self.traj) - cutoff_end):	×
61	if ind % steps == 0.0:	×
62	velocities_new = self.traj[ind].get_velocities()	×
63	velocities.append(velocities_new)	×
64
65	return np.array(velocities, dtype=np.float64)	×
66
67	def get_velocities_mass_weighted(	×
68	self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
69	) -> npt.NDArray[np.float64]:
70	"""
71	Weighs velocities by atomic masses.
72
73	Parameters
74	----------
75	steps : int, optional
76	Read every nth step. The default is 1 -> all steps are read. If for
77	instance steps=5 every 5th step is read.
78	cutoff_start : int, optional
79	Cutoff n stept at the beginning of the trajectory. The default is 0.
80	cutoff_end : int, optional
81	Cutoff n stept at the end of the trajectory. The default is 0.
82
83	Returns
84	-------
85	velocities_mass_weighted : np.array
86	Velocities weighted by atomic masses.
87
88	"""
89	velocities = self.get_velocities(	×
90	steps=steps, cutoff_start=cutoff_start, cutoff_end=cutoff_end
91	)
92
93	velocities_mass_weighted = np.zeros_like(velocities)	×
94	atomic_masses = self.geom.get_atomic_masses()	×
95
96	for i in range(velocities.shape[1]):	×
97	velocities_mass_weighted[:, i, :] = velocities[:, i, :] * np.sqrt(	×
98	atomic_masses[i]
99	)
100
101	return velocities_mass_weighted	×
102
103	def get_velocities_projected(	×
104	self,
105	projection_vectors: npt.NDArray[np.float64],
106	mass_weighted: bool = True,
107	steps: int = 1,
108	cutoff_start: int = 0,
109	cutoff_end: int = 0,
110	use_numba: bool = True,
111	) -> npt.NDArray[np.float64]:
112	"""
113	Calculate the normal-mode-decomposition of the velocities. This is done
114	by projecting the atomic velocities onto the vibrational eigenvectors.
115	See equation 10 in: https://doi.org/10.1016/j.cpc.2017.08.017.
116
117	Parameters
118	----------
119	projection_vectors : npt.NDArray[np.float64]
120	Array containing the vectors onto which the velocities should be
121	projected. Normally you will want this to be the eigenvectors
122	return by dfttoolkit.vibration.get_eigenvalues_and_eigenvectors.
123	mass_weighted : bool
124	Wheter the velocities should be mass weighted. If you do
125	normal-mode-decomposition with eigenvectors than this should be
126	True. The default is True.
127	steps : int, optional
128	Read every nth step. The default is 1 -> all steps are read. If for
129	instance steps=5 every 5th step is read.
130	cutoff_start : int, optional
131	Cutoff n stept at the beginning of the trajectory. The default is 0.
132	cutoff_end : int, optional
133	Cutoff n stept at the end of the trajectory. The default is 0.
134	use_numba : bool
135	Use numba for projection. Usually faster, but if your numbe
136	installation is wonky, it may produce falty results.
137
138	Returns
139	-------
140	velocities_projected : npt.NDArray[np.float64]
141	Velocities projected onto the eigenvectors structured as follows:
142	[number of time steps, number of frequencies]
143
144	"""
145	if mass_weighted:	×
146	velocities = self.get_velocities_mass_weighted(	×
147	steps=steps, cutoff_start=cutoff_start, cutoff_end=cutoff_end
148	)
149	else:
150	velocities = self.get_velocities(	×
151	steps=steps, cutoff_start=cutoff_start, cutoff_end=cutoff_end
152	)
153
154	return vu.get_normal_mode_decomposition(	×
155	velocities,
156	projection_vectors,
157	use_numba=use_numba,
158	)
159
160
161	def get_kinetic_energies(	×
162	self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
163	) -> npt.NDArray[np.float64]:
164	"""
165	Weighs velocities by atomic masses.
166
167	Parameters
168	----------
169	steps : int, optional
170	Read every nth step. The default is 1 -> all steps are read. If for
171	instance steps=5 every 5th step is read.
172	cutoff_start : int, optional
173	Cutoff n stept at the beginning of the trajectory. The default is 0.
174	cutoff_end : int, optional
175	Cutoff n stept at the end of the trajectory. The default is 0.
176
177	Returns
178	-------
179	velocities_mass_weighted : np.array
180	Velocities weighted by atomic masses.
181
182	"""
183	velocities = self.get_velocities_mass_weighted(	×
184	steps=steps, cutoff_start=cutoff_start, cutoff_end=cutoff_end
185	)
186	velocities_norm = np.linalg.norm(velocities, axis=2)	×
187
188	return 0.5 * velocities_norm**2	×
189
190
191	def get_temperature(	×
192	self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
193	) -> npt.NDArray[np.float64]:
194	"""
195	Get atomic temperature along an MD trajectory.
196
197	Parameters
198	----------
199	steps : int, optional
200	Read every nth step. The default is 1 -> all steps are read. If for
201	instance steps=5 every 5th step is read.
202	cutoff_start : int, optional
203	Cutoff n stept at the beginning of the trajectory. The default is 0.
204	cutoff_end : int, optional
205	Cutoff n stept at the end of the trajectory. The default is 0.
206
207	Returns
208	-------
209	temperature : npt.NDArray[np.float64]
210
211	"""
212	temperature = []	×
213
214	for ind in range(cutoff_start, len(self.traj) - cutoff_end):	×
215	if ind % steps == 0.0:	×
216	atoms = self.traj[ind]	×
217	unconstrained_atoms = len(atoms) - len(atoms.constraints[0].index)	×
218
219	ekin = atoms.get_kinetic_energy() / unconstrained_atoms	×
220	temperature.append(ekin / (1.5 * units.kB))	×
221
222	return np.array(temperature, dtype=np.float64)	×
223
224	def get_total_energy(	×
225	self, steps: int = 1, cutoff_start: int = 0, cutoff_end: int = 0
226	) -> npt.NDArray[np.float64]:
227	"""
228	Get atomic total energy along an MD trajectory.
229
230	Parameters
231	----------
232	steps : int, optional
233	Read every nth step. The default is 1 -> all steps are read. If for
234	instance steps=5 every 5th step is read.
235	cutoff_start : int, optional
236	Cutoff n stept at the beginning of the trajectory. The default is 0.
237	cutoff_end : int, optional
238	Cutoff n stept at the end of the trajectory. The default is 0.
239
240	Returns
241	-------
242	total_energy : npt.NDArray[np.float64]
243
244	"""
245	total_energy = []	×
246
247	for ind in range(cutoff_start, len(self.traj) - cutoff_end):	×
248	if ind % steps == 0.0:	×
249	atoms = self.traj[ind]	×
250	total_energy.append(atoms.get_total_energy())	×
251
252	return np.array(total_energy, dtype=np.float64)	×
253
254	def get_coords(self, atoms):	×
255	unconstrained_atoms = atoms.constraints[0].index	×
256
257	coords = []	×
258	for ind, atom in enumerate(atoms):	×
259	if ind not in unconstrained_atoms:	×
260	coords.append(atom.position)	×
261
262	return np.array(coords)	×
263
264	def get_atomic_displacements(	×
265	self,
266	coords_0=None,
267	steps: int = 1,
268	cutoff_start: int = 0,
269	cutoff_end: int = 0,
270	) -> npt.NDArray[np.float64]:
271	"""
272	Get atomic atomic displacements with respect to the first time step
273	along an MD trajectory.
274
275	Parameters
276	----------
277	steps : int, optional
278	Read every nth step. The default is 1 -> all steps are read. If for
279	instance steps=5 every 5th step is read.
280	cutoff_start : int, optional
281	Cutoff n stept at the beginning of the trajectory. The default is 0.
282	cutoff_end : int, optional
283	Cutoff n stept at the end of the trajectory. The default is 0.
284
285	Returns
286	-------
287	atomic_displacements : npt.NDArray[np.float64]
288
289	"""
290	atomic_displacements = []	×
291
292	if coords_0 is None:	×
293	coords_0 = self.get_coords(self.traj[cutoff_start])	×
294
295	for ind in range(cutoff_start, len(self.traj) - cutoff_end):	×
296	if ind % steps == 0.0:	×
297	coords = self.get_coords(self.traj[ind])	×
298
299	disp = np.linalg.norm(coords - coords_0, axis=1)	×
300	atomic_displacements.append(disp)	×
301
302	return np.array(atomic_displacements, dtype=np.float64)	×

maurergroup / dfttoolkit / 15065810745

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