.. note::
    :class: sphx-glr-download-link-note

    Click :ref:`here <sphx_glr_download_code_vib.py>` to download the full example code
.. rst-class:: sphx-glr-example-title

.. _sphx_glr_code_vib.py:


Vibrational Analysis
====================

The module ``nnp.vib`` contains tools to compute analytical hessian
and do vibrational analysis.

Let's first import all the packages we will use:


.. code-block:: default

    import torch
    from torch import Tensor
    from typing import NamedTuple, Optional









With ``torch.autograd`` automatic differentiation engine, computing analytical
hessian is very simple: Just compute the gradient of energies with respect
to forces, and then compute the gradient of each element of forces with respect
to coordinates again.

The ``torch.autograd.grad`` returns a list of ``Optional[Tensor]``. To be
compatible with ``torch.jit``, we need to use assert statement to allow
``torch.jit`` to do `type refinement`_.

.. _type refinement:
  https://pytorch.org/docs/stable/jit.html#optional-type-refinement


.. code-block:: default

    def _get_derivatives_not_none(x: Tensor, y: Tensor, retain_graph: Optional[bool] = None,
                                  create_graph: bool = False) -> Tensor:
        ret = torch.autograd.grad(
            [y.sum()], [x], retain_graph=retain_graph, create_graph=create_graph)[0]
        assert ret is not None
        return ret


    def hessian(coordinates: Tensor, energies: Optional[Tensor] = None,
                forces: Optional[Tensor] = None) -> Tensor:
        """Compute analytical hessian from the energy graph or force graph.

        Arguments:
            coordinates: Tensor of shape `(molecules, atoms, 3)` or `(atoms, 3)`
            energies: Tensor of shape `(molecules,)`, or scalar, if specified,
                then `forces` must be `None`. This energies must be computed
                from `coordinates` in a graph.
            forces: Tensor of shape `(molecules, atoms, 3)` or `(atoms, 3)`,
                if specified, then `energies` must be `None`. This forces must
                be computed from `coordinates` in a graph.

        Returns:
            Tensor of shape `(molecules, 3 * atoms, 3 * atoms)` or `(3 * atoms, 3 * atoms)`
        """
        if energies is None and forces is None:
            raise ValueError('Energies or forces must be specified')
        if energies is not None and forces is not None:
            raise ValueError('Energies or forces can not be specified at the same time')
        if forces is None:
            assert energies is not None
            forces = -_get_derivatives_not_none(coordinates, energies, create_graph=True)
        flattened_force = forces.flatten(start_dim=-2)
        force_components = flattened_force.unbind(dim=-1)
        return -torch.stack([
            _get_derivatives_not_none(coordinates, f, retain_graph=True).flatten(start_dim=-2)
            for f in force_components
        ], dim=-1)









Below are helper functions to compute vibrational frequencies and normal modes.
The normal modes and vibrational frquencies satisfies the following equation.

.. math::
  H q = \omega^2 T q

where :math:`H` is the Hessian matrix, :math:`q` is the normal coordinates, and

.. math::
  T=\left[
   \begin{array}{ccccccc}
   m_{1}\\
   & m_{1}\\
   &  & m_{1}\\
   &  &  & m_{2}\\
   &  &  &  & m_{2}\\
   &  &  &  &  & m_{2}\\
   &  &  &  &  &  & \ddots
   \end{array}
   \right]

is the mass for each coordinate.

This is a generalized eigen problem, which is not immediately supported by PyTorch
So we solve this problem through Lowdin diagnolization:

.. math::
  H q = \omega^2 T q \Longrightarrow H q = \omega^2 T^{\frac{1}{2}} T^{\frac{1}{2}} q

Let :math:`q' = T^{\frac{1}{2}} q`, we then have

.. math::
  T^{\frac{1}{2}} H T^{\frac{1}{2}} q' = \omega^2 q'

this is a regular eigen problem


.. code-block:: default

    class FreqsModes(NamedTuple):
        angular_frequencies: Tensor
        modes: Tensor


    def vibrational_analysis(masses: Tensor, hessian: Tensor) -> FreqsModes:
        """Computing the vibrational wavenumbers from hessian.

        Arguments:
            masses: Tensor of shape `(molecules, atoms)` or `(atoms,)`.
            hessian: Tensor of shape `(molecules, 3 * atoms, 3 * atoms)` or
                `(3 * atoms, 3 * atoms)`.

        Returns:
            A namedtuple `(angular_frequencies, modes)` where

            angular_frequencies:
                Tensor of shape `(molecules, 3 * atoms)` or `(3 * atoms,)`
            modes:
                Tensor of shape `(molecules, modes, atoms, 3)` or `(modes, atoms, 3)`
                where `modes = 3 * atoms` is the number of normal modes.
        """
        inv_sqrt_mass = masses.rsqrt().repeat_interleave(3, dim=-1)
        mass_scaled_hessian = hessian * inv_sqrt_mass.unsqueeze(-2) * inv_sqrt_mass.unsqueeze(-1)
        eigenvalues, eigenvectors = torch.symeig(mass_scaled_hessian, eigenvectors=True)
        angular_frequencies = eigenvalues.sqrt()
        modes = (eigenvectors.transpose(-1, -2) * inv_sqrt_mass.unsqueeze(-2))
        new_shape = modes.shape[:-1] + (-1, 3)
        modes = modes.reshape(new_shape)
        return FreqsModes(angular_frequencies, modes)








.. rst-class:: sphx-glr-timing

   **Total running time of the script:** ( 0 minutes  0.002 seconds)


.. _sphx_glr_download_code_vib.py:


.. only :: html

 .. container:: sphx-glr-footer
    :class: sphx-glr-footer-example



  .. container:: sphx-glr-download

     :download:`Download Python source code: vib.py <vib.py>`



  .. container:: sphx-glr-download

     :download:`Download Jupyter notebook: vib.ipynb <vib.ipynb>`


.. only:: html

 .. rst-class:: sphx-glr-signature

    `Gallery generated by Sphinx-Gallery <https://sphinx-gallery.github.io>`_