.. only:: html

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

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

    .. _sphx_glr_code_md.py:


Molecular Dynamics
==================

The module ``nnp.md`` provide tools to run molecular dynamics with a potential
defined by PyTorch.








.. code-block:: default


    import torch
    from torch import Tensor
    from nnp import pbc
    from typing import Callable, Sequence
    import ase.calculators.calculator


    class Calculator(ase.calculators.calculator.Calculator):
        """ASE Calculator that wraps a neural network potential

        Arguments:
            func (callable): A fucntion that .
            overwrite (bool): After wrapping atoms into central box, whether
                to replace the original positions stored in :class:`ase.Atoms`
                object with the wrapped positions.
        """

        implemented_properties = ['energy', 'forces', 'stress', 'free_energy']

        def __init__(self, func: Callable[[Sequence[str], Tensor, Tensor, Tensor], Tensor],
                     overwrite: bool = False):
            super(Calculator, self).__init__()
            self.func = func
            self.overwrite = overwrite

        def calculate(self, atoms=None, properties=['energy'],
                      system_changes=ase.calculators.calculator.all_changes):
            super(Calculator, self).calculate(atoms, properties, system_changes)
            coordinates = torch.from_numpy(self.atoms.get_positions()).requires_grad_('forces' in properties)
            cell = coordinates.new_tensor(self.atoms.get_cell(complete=True).array)
            pbc_ = torch.tensor(self.atoms.get_pbc(), dtype=torch.bool)
            pbc_enabled = pbc_.any().item()

            if pbc_enabled:
                coordinates = pbc.map2central(cell, coordinates, pbc_)

            if 'stress' in properties:
                scaling = torch.eye(3, requires_grad=True)
                coordinates = coordinates @ scaling
                cell = cell @ scaling

            energy = self.func(atoms.get_chemical_symbols(), coordinates, cell, pbc_)

            self.results['energy'] = energy.item()
            self.results['free_energy'] = energy.item()

            if 'forces' in properties:
                forces = -torch.autograd.grad(energy, coordinates)[0]
                self.results['forces'] = forces.cpu().numpy()

            if 'stress' in properties:
                volume = self.atoms.get_volume()
                stress = torch.autograd.grad(energy, scaling)[0] / volume
                self.results['stress'] = stress.cpu().numpy()


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

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


.. _sphx_glr_download_code_md.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: md.py <md.py>`



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

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


.. only:: html

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

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