.. only:: html

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

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

    .. _sphx_glr_examples_planetary_motion.py:


Use ASE to Simulate Planetary Motion
====================================

This tutorial shows how to use `ASE`_ to simulate planetary motion.
Although both ASE and NNP are designed for molecular dynamics (MD),
implementing a force field to do MD is not at a reasonable abount of
work to use as a tutorial. So we simulate something different here.
Assume we have two atoms, one hydrogen and one carbon. They are all
moving under the gravity centered at the origin. There is no interaction
between these two atoms.

.. _ASE:
    https://wiki.fysik.dtu.dk/ase

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


.. code-block:: default

    import ase
    from ase.data import atomic_masses
    from ase.units import fs
    from ase.md.verlet import VelocityVerlet
    import nnp.md as md
    import pytest
    import sys
    import math
    import matplotlib.pyplot as plt









The module `nnp.md` provide tools to run molecular dynamics with a potential
defined by PyTorch. To define the potential we need to provide a function that
takes the chemical symbols, atom coordinates, unit cell, and PBC flag as an input
and returns the molecular energy as an output. The output and input coordinates
and cell must be in the a differentiable graph so that PyTorch could automatically
compute the force and stress tensor using the `autograd` engine. In this example,
we don't care about the periodic boundary condition and unit cell because we
are actually not simulating molecular system.

From basic physics, we know that the equation for gravitational potential has
the equation :math:`E(r)=-\frac{GMm}{r}`. For simplicity, we choose :math:`GM=1`
when :math:`E` is in :math:`\mathrm{eV}`, :math:`m` in :math:`\mathrm{amu}`,
and :math:`r` in Angstrom.


.. code-block:: default

    def gravity(chemical_symbols, coordinates, _cell, _pbc):
        assert len(chemical_symbols) == coordinates.shape[0]
        atomic_numbers = {'H': 0, 'C': 6}
        mass = coordinates.new_tensor(
            [atomic_masses[atomic_numbers[s]] for s in chemical_symbols])
        inv_r = 1 / coordinates.norm(dim=1)
        potential_energies = - mass * inv_r
        return potential_energies.sum()









with this potential, we could define an ASE calculator:


.. code-block:: default

    calculator = md.Calculator(gravity)








Now let's create two atoms, initially at (1, 0, 0) and (2, 0, 0)


.. code-block:: default

    planets = ase.Atoms('CH', [[1, 0, 0], [2, 0, 0]])








For the purpose of demonstration, we make these two atoms's trajector a perfect
circle at the XY plane. To do so, we need to carefully set the initial velocity.
It is not hard to get that, the velocity is :math:`v = \sqrt{\frac{GM}{r}}`.


.. code-block:: default

    planets.set_velocities([[0, 1, 0], [0, 1 / math.sqrt(2), 0]])
    planets.set_calculator(calculator)








Now we can start the dynamics:


.. code-block:: default

    X1 = []
    Y1 = []
    Z1 = []
    X2 = []
    Y2 = []
    Z2 = []


    def dump_positions():
        (x1, y1, z1), (x2, y2, z2) = planets.get_positions()
        X1.append(x1)
        Y1.append(y1)
        Z1.append(z1)
        X2.append(x2)
        Y2.append(y2)
        Z2.append(z2)


    dyn = VelocityVerlet(planets, timestep=0.01 * fs)
    dyn.attach(dump_positions)
    dyn.run(20000)





.. rst-class:: sphx-glr-script-out

 Out:

 .. code-block:: none


    True



Now let's plot the trajectory to see what we get:


.. code-block:: default

    if __name__ == '__main__':
        plt.clf()
        plt.plot(X1, Y1, label='C')
        plt.plot(X2, Y2, label='H')
        plt.legend()
        plt.axes().set_aspect('equal')
        plt.show()





.. image:: /examples/images/sphx_glr_planetary_motion_001.png
    :alt: planetary motion
    :class: sphx-glr-single-img


.. rst-class:: sphx-glr-script-out

 Out:

 .. code-block:: none

    /home/runner/work/nnp/nnp/tests/planetary_motion.py:100: MatplotlibDeprecationWarning: Adding an axes using the same arguments as a previous axes currently reuses the earlier instance.  In a future version, a new instance will always be created and returned.  Meanwhile, this warning can be suppressed, and the future behavior ensured, by passing a unique label to each axes instance.
      plt.axes().set_aspect('equal')




To check the correctness of this dynamics, we use pytest to test if it is on
the XY plane and if the trajectory is circle:


.. code-block:: default

    def test_is_on_xy_plan():
        for z in Z1 + Z2:
            assert abs(z) < 1e-4


    def test_is_circle():
        for x, y in zip(X1, Y1):
            assert abs(math.sqrt(x * x + y * y) - 1) < 0.1
        for x, y in zip(X2, Y2):
            assert abs(math.sqrt(x * x + y * y) - 2) < 0.1


    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 2 items

    planetary_motion.py::test_is_on_xy_plan PASSED                           [ 50%]
    planetary_motion.py::test_is_circle PASSED                               [100%]

    ============================== 2 passed in 24.56s ==============================





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

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


.. _sphx_glr_download_examples_planetary_motion.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: planetary_motion.py <planetary_motion.py>`



  .. container:: sphx-glr-download sphx-glr-download-jupyter

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


.. only:: html

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

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