.. only:: html
.. note::
:class: sphx-glr-download-link-note
Click :ref:`here ` to download the full example code
.. rst-class:: sphx-glr-example-title
.. _sphx_glr_examples_vibration.py:
Analytical Hessians and Vibrational Analysis
============================================
This tutorial demonstrates how compute analytical hessians and do
vibrational analysis using ``nnp.vib``.
Let's first import all the packages we will use:
.. code-block:: default
import torch
import math
import pytest
from pytest import approx
import sys
import nnp.so3 as so3
import nnp.vib as vib
In this tutorial, we will study atoms moving in a quadratic potential.
There is no interaction between these atoms.
Let's first construct such a potential. A naive potential would be:
.. math::
U(x, y, z) = \frac{1}{2} \left(0.5 x^2 + y^2 + 2 z^2\right)
.. code-block:: default
def naive_potential(x, y, z):
return 0.5 * (0.5 * x ** 2 + y ** 2 + 2 * z ** 2)
A naive potential is not very interesting, let's rotate the potential along
the z axis for 45 degrees. Rotating the potential along the z axis for 45 degrees
is equivalent to rotate the coordinates along the z axis for -45 degrees before
evaluating the naive potential. Let's make our potential able to handle both
single molecule (i.e. coordinates has shape `(atoms, 3)`), and in batch (i.e.
coordinates has shape `(molecules, atoms, 3)`):
.. code-block:: default
rot45 = so3.rotate_along(torch.tensor([0, 0, math.pi / 4]))
rot_neg_45 = so3.rotate_along(torch.tensor([0, 0, -math.pi / 4]))
def potential(coordinates):
rotated_coordinates = (rot_neg_45 @ coordinates.transpose(-1, -2)).transpose(-1, -2)
x, y, z = rotated_coordinates.unbind(-1)
return naive_potential(x, y, z).sum(dim=-1)
Analytical Hessian
------------------
Now let's compute the hessian one molecule containing two atoms at the origin
.. code-block:: default
coordinates = torch.zeros(2, 3, requires_grad=True)
energy = potential(coordinates)
hessian = vib.hessian(coordinates, energies=energy)
print(hessian)
.. rst-class:: sphx-glr-script-out
Out:
.. code-block:: none
tensor([[ 0.7500, -0.2500, -0.0000, -0.0000, -0.0000, -0.0000],
[-0.2500, 0.7500, -0.0000, -0.0000, -0.0000, -0.0000],
[-0.0000, -0.0000, 2.0000, -0.0000, -0.0000, -0.0000],
[-0.0000, -0.0000, -0.0000, 0.7500, -0.2500, -0.0000],
[-0.0000, -0.0000, -0.0000, -0.2500, 0.7500, -0.0000],
[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000, 2.0000]])
Let's compute the theoretical result to see if it matches with the result above:
First of all, because there are no interactions between atoms, the hessian
should be a block diagonal matrix:
.. math::
H = \left[9\times9\right] = \left[\begin{array}{cc}
3\times3\\
& 3\times3
\end{array}\right]
The two :math:`3\times3` matricies are identical. For each :math:`3\times3`
matrix, the potential is not rotated at the z axis, so the structure of it
should be:
.. math::
\left[\begin{array}{ccc}
? & ?\\
? & ?\\
& & 2
\end{array}\right]
It is not hard to figure out that, for the rotated potential, considering
only the contribution from x,y plane, it can be written as:
.. math::
U=\frac{1}{2}\left[0.5\left(\frac{x+y}{\sqrt{2}}\right)^{2}+\left(\frac{y-x}{\sqrt{2}}\right)^{2}\right]=\frac{1}{4}\left(1.5x^{2}+1.5y^{2}-xy\right)
Therefore, the :math:`2\times2` block on the top left should be
.. math::
\left[\begin{array}{cc}
0.75 & -0.25\\
-0.25 & 0.75
\end{array}\right]
.. code-block:: default
def test_analytical_hessian():
hessian00 = hessian[:3, :3]
hessian01 = hessian[:3, 3:]
hessian10 = hessian[3:, :3]
hessian11 = hessian[3:, 3:]
expected = torch.tensor([
[ 0.75, -0.25, 0], # noqa: E201, E241
[-0.25, 0.75, 0], # noqa: E201, E241
[ 0.00, 0.00, 2], # noqa: E201, E241
])
assert torch.allclose(hessian00, expected)
assert torch.allclose(hessian11, expected)
assert torch.allclose(hessian10, torch.zeros(3, 3))
assert torch.allclose(hessian01, torch.zeros(3, 3))
We also support compute multiple molecules in batch
.. code-block:: default
coordinates_batch = torch.zeros(2, 2, 3, requires_grad=True)
energy_batch = potential(coordinates_batch)
hessian_batch = vib.hessian(coordinates_batch, energies=energy_batch)
print(hessian_batch)
.. rst-class:: sphx-glr-script-out
Out:
.. code-block:: none
tensor([[[ 0.7500, -0.2500, -0.0000, -0.0000, -0.0000, -0.0000],
[-0.2500, 0.7500, -0.0000, -0.0000, -0.0000, -0.0000],
[-0.0000, -0.0000, 2.0000, -0.0000, -0.0000, -0.0000],
[-0.0000, -0.0000, -0.0000, 0.7500, -0.2500, -0.0000],
[-0.0000, -0.0000, -0.0000, -0.2500, 0.7500, -0.0000],
[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000, 2.0000]],
[[ 0.7500, -0.2500, -0.0000, -0.0000, -0.0000, -0.0000],
[-0.2500, 0.7500, -0.0000, -0.0000, -0.0000, -0.0000],
[-0.0000, -0.0000, 2.0000, -0.0000, -0.0000, -0.0000],
[-0.0000, -0.0000, -0.0000, 0.7500, -0.2500, -0.0000],
[-0.0000, -0.0000, -0.0000, -0.2500, 0.7500, -0.0000],
[-0.0000, -0.0000, -0.0000, -0.0000, -0.0000, 2.0000]]])
The hessian should be just the stack of the previous results twice
.. code-block:: default
def test_analytical_hessian_batch():
expected = torch.stack([hessian, hessian])
assert torch.allclose(expected, hessian_batch)
Vibrational Analysis
--------------------
Now let's do vibrational analysis using the computed hessians. Let's assume
the two atoms has mass (1, 3). The hessian should have shape `(3 * atoms, 3 * atoms)`,
the mass should have shape `(atoms,)`
.. code-block:: default
mass = torch.tensor([1.0, 3.0])
freq_modes1 = vib.vibrational_analysis(mass, hessian)
The output angular frequencies should have shape `(modes,)`, while the modes
have shape `(modes, atoms, 3)`, where the modes dimension correspond to different
vibrational modes
.. code-block:: default
print(freq_modes1.angular_frequencies.shape)
print(freq_modes1.angular_frequencies)
print(freq_modes1.modes.shape)
print(freq_modes1.modes)
.. rst-class:: sphx-glr-script-out
Out:
.. code-block:: none
torch.Size([6])
tensor([0.4082, 0.5774, 0.7071, 0.8165, 1.0000, 1.4142])
torch.Size([6, 2, 3])
tensor([[[-0.0000, -0.0000, -0.0000],
[-0.4082, -0.4082, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[-0.4082, 0.4082, 0.0000]],
[[-0.7071, -0.7071, -0.0000],
[-0.0000, -0.0000, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[ 0.0000, 0.0000, 0.5774]],
[[-0.7071, 0.7071, 0.0000],
[ 0.0000, 0.0000, 0.0000]],
[[ 0.0000, 0.0000, 1.0000],
[ 0.0000, 0.0000, 0.0000]]])
The angular frequency of a harmonic oscillator is given by :math:`\sqrt{\frac{k}{m}}`.
For this example, the two atoms are independent, therefore the normal modes would
have one atom moving while the other static. The angular frequencies are also
independent to each other. Therefore, we expect to see the following combinations
of angular frequencies and normal modes:
+----------------------------+-------------+-----------+
| Angular Frequency | Moving Atom | Direction |
+============================+=============+===========+
| :math:`\frac{1}{\sqrt{2}}` | 1 | (x, y) |
+----------------------------+-------------+-----------+
| :math:`1` | 1 | (x, -y) |
+----------------------------+-------------+-----------+
| :math:`\sqrt{2}` | 1 | z |
+----------------------------+-------------+-----------+
| :math:`\frac{1}{\sqrt{6}}` | 2 | (x, y) |
+----------------------------+-------------+-----------+
| :math:`\frac{1}{\sqrt{3}}` | 2 | (x, -y) |
+----------------------------+-------------+-----------+
| :math:`\sqrt{\frac{2}{3}}` | 2 | z |
+----------------------------+-------------+-----------+
Now let's write a test function to test this case
.. code-block:: default
def test_vibrational_analysis():
expected_freqs = torch.tensor([
1 / math.sqrt(6), 1 / math.sqrt(3), 1 / math.sqrt(2),
math.sqrt(2 / 3), 1, math.sqrt(2)])
assert torch.allclose(freq_modes1.angular_frequencies, expected_freqs)
mode0 = freq_modes1.modes[0]
atom0, atom1 = mode0
assert torch.allclose(atom0, torch.zeros(3))
assert atom1[0].item() == approx(atom1[1].item())
assert atom1[2].item() == approx(0)
mode1 = freq_modes1.modes[1]
atom0, atom1 = mode1
assert torch.allclose(atom0, torch.zeros(3))
assert atom1[0].item() == approx(-atom1[1].item())
assert atom1[2].item() == approx(0)
mode2 = freq_modes1.modes[2]
atom0, atom1 = mode2
assert torch.allclose(atom1, torch.zeros(3))
assert atom0[0].item() == approx(atom0[1].item())
assert atom0[2].item() == approx(0)
mode3 = freq_modes1.modes[3]
atom0, atom1 = mode3
assert torch.allclose(atom0, torch.zeros(3))
assert torch.allclose(atom1[:1], torch.zeros(2))
assert atom1[2].abs().item() > 1e-3
mode4 = freq_modes1.modes[4]
atom0, atom1 = mode4
assert torch.allclose(atom1, torch.zeros(3))
assert atom0[0].item() == approx(-atom0[1].item())
assert atom0[2].item() == approx(0)
mode5 = freq_modes1.modes[5]
atom0, atom1 = mode5
assert torch.allclose(atom1, torch.zeros(3))
assert torch.allclose(atom0[:1], torch.zeros(2))
assert atom0[2].abs().item() > 1e-3
We also support doing vibrational analysis in batch. In case of batch, the
hessian should have shape `(molecules, 3 * atoms, 3 * atoms)`, and the mass
should have shape `(molecules, atoms)`. The resulting tensors is just a stack
of tensors for the single molecule case.
.. code-block:: default
mass_batch = torch.stack([mass, mass])
freq_modes2 = vib.vibrational_analysis(mass_batch, hessian_batch)
print(freq_modes2.angular_frequencies.shape)
print(freq_modes2.angular_frequencies)
print(freq_modes2.modes.shape)
print(freq_modes2.modes)
.. rst-class:: sphx-glr-script-out
Out:
.. code-block:: none
torch.Size([2, 6])
tensor([[0.4082, 0.5774, 0.7071, 0.8165, 1.0000, 1.4142],
[0.4082, 0.5774, 0.7071, 0.8165, 1.0000, 1.4142]])
torch.Size([2, 6, 2, 3])
tensor([[[[-0.0000, -0.0000, -0.0000],
[-0.4082, -0.4082, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[-0.4082, 0.4082, 0.0000]],
[[-0.7071, -0.7071, -0.0000],
[-0.0000, -0.0000, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[ 0.0000, 0.0000, 0.5774]],
[[-0.7071, 0.7071, 0.0000],
[ 0.0000, 0.0000, 0.0000]],
[[ 0.0000, 0.0000, 1.0000],
[ 0.0000, 0.0000, 0.0000]]],
[[[-0.0000, -0.0000, -0.0000],
[-0.4082, -0.4082, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[-0.4082, 0.4082, 0.0000]],
[[-0.7071, -0.7071, -0.0000],
[-0.0000, -0.0000, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[ 0.0000, 0.0000, 0.5774]],
[[-0.7071, 0.7071, 0.0000],
[ 0.0000, 0.0000, 0.0000]],
[[ 0.0000, 0.0000, 1.0000],
[ 0.0000, 0.0000, 0.0000]]]])
Let's test if the result is a stack of pervious results
.. code-block:: default
def test_vibrational_analysis_batch():
af1 = freq_modes1.angular_frequencies
m1 = freq_modes1.modes
expected_af = torch.stack([af1, af1])
expected_modes = torch.stack([m1, m1])
assert torch.allclose(expected_af, freq_modes2.angular_frequencies)
assert torch.allclose(expected_modes, freq_modes2.modes)
In the case when the batch are isomers (i.e. have the same masses), it is also
possible to just specify the mass as shape `(atoms,)` and it will broadcast
automatically.
.. code-block:: default
freq_modes3 = vib.vibrational_analysis(mass, hessian_batch)
print(freq_modes3.angular_frequencies.shape)
print(freq_modes3.angular_frequencies)
print(freq_modes3.modes.shape)
print(freq_modes3.modes)
.. rst-class:: sphx-glr-script-out
Out:
.. code-block:: none
torch.Size([2, 6])
tensor([[0.4082, 0.5774, 0.7071, 0.8165, 1.0000, 1.4142],
[0.4082, 0.5774, 0.7071, 0.8165, 1.0000, 1.4142]])
torch.Size([2, 6, 2, 3])
tensor([[[[-0.0000, -0.0000, -0.0000],
[-0.4082, -0.4082, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[-0.4082, 0.4082, 0.0000]],
[[-0.7071, -0.7071, -0.0000],
[-0.0000, -0.0000, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[ 0.0000, 0.0000, 0.5774]],
[[-0.7071, 0.7071, 0.0000],
[ 0.0000, 0.0000, 0.0000]],
[[ 0.0000, 0.0000, 1.0000],
[ 0.0000, 0.0000, 0.0000]]],
[[[-0.0000, -0.0000, -0.0000],
[-0.4082, -0.4082, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[-0.4082, 0.4082, 0.0000]],
[[-0.7071, -0.7071, -0.0000],
[-0.0000, -0.0000, -0.0000]],
[[ 0.0000, 0.0000, 0.0000],
[ 0.0000, 0.0000, 0.5774]],
[[-0.7071, 0.7071, 0.0000],
[ 0.0000, 0.0000, 0.0000]],
[[ 0.0000, 0.0000, 1.0000],
[ 0.0000, 0.0000, 0.0000]]]])
Let's test if we still have the same result
.. code-block:: default
def test_vibrational_analysis_batch_broadcast():
assert torch.allclose(freq_modes2.angular_frequencies, freq_modes3.angular_frequencies)
assert torch.allclose(freq_modes2.modes, freq_modes3.modes)
Now let's run all the tests
.. code-block:: default
if __name__ == '__main__':
pytest.main([sys.argv[0], '-v'])
.. rst-class:: sphx-glr-script-out
Out:
.. code-block:: none
============================= test session starts ==============================
platform linux -- Python 3.7.7, pytest-5.4.3, py-1.8.2, pluggy-0.13.1 -- /opt/hostedtoolcache/Python/3.7.7/x64/bin/python
cachedir: .pytest_cache
rootdir: /home/runner/work/nnp/nnp, inifile: setup.cfg
collecting ... collected 5 items
vibration.py::test_analytical_hessian PASSED [ 20%]
vibration.py::test_analytical_hessian_batch PASSED [ 40%]
vibration.py::test_vibrational_analysis PASSED [ 60%]
vibration.py::test_vibrational_analysis_batch PASSED [ 80%]
vibration.py::test_vibrational_analysis_batch_broadcast PASSED [100%]
============================== 5 passed in 0.05s ===============================
.. rst-class:: sphx-glr-timing
**Total running time of the script:** ( 0 minutes 0.135 seconds)
.. _sphx_glr_download_examples_vibration.py:
.. only :: html
.. container:: sphx-glr-footer
:class: sphx-glr-footer-example
.. container:: sphx-glr-download sphx-glr-download-python
:download:`Download Python source code: vibration.py `
.. container:: sphx-glr-download sphx-glr-download-jupyter
:download:`Download Jupyter notebook: vibration.ipynb `
.. only:: html
.. rst-class:: sphx-glr-signature
`Gallery generated by Sphinx-Gallery `_