from typing import Optional
import numpy as np
from ase import Atoms
from numpy.typing import NDArray
from dynasor.logging_tools import logger
from dynasor.modes.atoms import DynasorAtoms
from dynasor.modes.tools import mode_coordinates, mode_forces, mode_momenta
from dynasor.trajectory import Trajectory
from dynasor.tools.structures import get_displacements_from_u
def _validated_masses(ideal_supercell: Atoms) -> NDArray[float]:
"""Returns the masses of the supercell in dmu, checking that every one of them can be
divided by."""
masses = DynasorAtoms(ideal_supercell).masses
bad = np.flatnonzero(~np.isfinite(masses) | (masses <= 0))
if len(bad) > 0:
raise ValueError(
'Projecting forces requires a positive and finite mass for every atom of '
f'ideal_supercell, but the masses of atoms {bad.tolist()} are not.')
return masses
[docs]
def project_modes(
traj: Trajectory,
modes: NDArray[float],
ideal_supercell: Atoms,
check_mic: Optional[bool] = True,
logging_interval: Optional[int] = 1000,
) -> tuple[NDArray[complex], NDArray[complex], NDArray[complex]]:
r"""Projects an atomic trajectory onto set of phonon modes.
The projection follows the conventions of :class:`ModeProjector
<dynasor.ModeProjector>`,
.. math::
Q = X u, \quad P = X^* v, \quad F = X^* f / m
where :math:`X` are the modes, :math:`u` the displacements, :math:`v` the velocities,
:math:`f` the forces and :math:`m` the atomic masses.
The mode forces are conjugate to the mode coordinates, meaning
:math:`F = -\partial V / \partial Q`, so the virial energy of a mode is :math:`-QF/2`.
See :attr:`ModeProjector.virial_energies <dynasor.ModeProjector.virial_energies>` for that
quantity and for how its imaginary parts behave.
Parameters
----------
traj
Input trajectory.
modes
Modes to project on, as an array with shape ``(..., N, 3)`` where ``N`` is the
number of atoms in the supercell and the leading dimensions define the output shape.
ideal_supercell
Ideal supercell used to find atomic displacements and to provide the atomic masses.
The masses have to be the ones that :attr:`modes` was built with, which for modes from
a :class:`ModeProjector <dynasor.ModeProjector>` are the supercell masses of that
projector.
Neither the atom ordering nor the masses are checked, and a mismatch changes `F` per
atom instead of by an overall factor.
The atom count is checked against :attr:`traj`, and a mismatched cell only triggers a
warning, since some thermal expansion relative to the reference structure is normal.
check_mic
Whether to wrap the displacements or not, faster if no wrap.
logging_interval
Log progress at ``INFO`` level every this many frames. Set to ``0`` to disable
progress logging.
Returns
-------
A tuple comprising `(Q, P, F)` where `Q` are the mode coordinates in Å√dmu, `P` are
the mode momenta in √eV and `F` are the mode forces in eV/Å√dmu.
Each array has shape ``(n_frames, *modes.shape[:-2])``, where ``n_frames`` is the
number of frames read from :attr:`traj`.
`P` is all zeros if :attr:`traj` provides no velocities, and `F` is all zeros if it
provides no forces.
Raises
------
ValueError
If :attr:`traj` provides no positions, if :attr:`modes` does not have shape
``(..., N, 3)``, if :attr:`ideal_supercell` and :attr:`traj` disagree on the number of
atoms, if a frame past the first provides no positions, velocities or forces, or if
:attr:`traj` provides forces while :attr:`ideal_supercell` gives an atom a mass that
is not positive and finite.
Examples
--------
The mode forces give the virial energy of every mode along a trajectory::
>>> Q, P, F = project_modes(traj, mp.eigenmodes, mp.supercell.to_ase()) # doctest: +SKIP
>>> virial_energies = (-Q * F / 2).real # doctest: +SKIP
"""
# logger
logger.info('Running mode projection')
if not traj.has_positions:
raise ValueError('project_modes requires positions to be available in the trajectory, '
'but traj does not provide positions.')
has_velocities = traj.has_velocities
has_forces = traj.has_forces
if not has_velocities:
logger.info('traj does not provide velocities; P will be returned as zeros')
if not has_forces:
logger.info('traj does not provide forces; F will be returned as zeros')
modes = np.asarray(modes)
if modes.ndim < 2 or modes.shape[-1] != 3 or modes.shape[-2] != traj.n_atoms:
raise ValueError(
f'modes must have shape (..., N, 3), where N = {traj.n_atoms} is the number of '
f'atoms in the trajectory, but has shape {modes.shape}.')
if traj.n_atoms != len(ideal_supercell):
raise ValueError('ideal_supercell must contain the same number of atoms as the trajectory.')
if not np.allclose(traj.cell, ideal_supercell.cell, atol=1e-5, rtol=0.0):
logger.warning('ideal_supercell cell does not match the trajectory cell.')
if has_forces:
masses = _validated_masses(ideal_supercell)
Q_traj, P_traj, F_traj = [], [], []
for it, frame in enumerate(traj):
if logging_interval and it % logging_interval == 0:
traj._log_progress(frame.frame_index)
# Make positions into displacements
x = frame.get_positions_as_array(traj._atomic_indices)
u = x - ideal_supercell.positions
# Calculate Q
u = get_displacements_from_u(u, ideal_supercell.cell, check_mic=check_mic)
Q_traj.append(mode_coordinates(modes, u))
# Calculate P
if has_velocities:
v = frame.get_velocities_as_array(traj._atomic_indices)
P_traj.append(mode_momenta(modes, v))
# Calculate F
if has_forces:
f = frame.get_forces_as_array(traj._atomic_indices)
F_traj.append(mode_forces(modes, f, masses))
# The reshape only acts on a trajectory of no frames, which would otherwise lose the
# leading mode dimensions.
Q_traj = np.array(Q_traj).reshape((-1, *modes.shape[:-2]))
P_traj = np.array(P_traj).reshape(Q_traj.shape) if has_velocities else np.zeros_like(Q_traj)
F_traj = np.array(F_traj).reshape(Q_traj.shape) if has_forces else np.zeros_like(Q_traj)
return Q_traj, P_traj, F_traj
