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

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

.. _sphx_glr_code_pbc.py:


Periodic Boundary Conditions
============================

The module ``nnp.pbc`` contains tools to deal with periodic boundary conditions.

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


.. code-block:: default

    import torch
    from torch import Tensor









The following function computes the number of repeats required to capture
all the neighbor atoms within the cutoff radius.


.. code-block:: default

    def num_repeats(cell: Tensor, pbc: Tensor, cutoff: float) -> Tensor:
        """Compute the number of repeats required along each cell vector to make
        the original cell and repeated cells together form a large enough box
        so that for each atom in the original cell, all its neighbor atoms within
        the given cutoff distance are contained in the big box.

        Arguments:
            cell: tensor of shape ``(3, 3)`` of the three vectors defining unit cell:

                .. code-block:: python

                    tensor([[x1, y1, z1],
                            [x2, y2, z2],
                            [x3, y3, z3]])

            pbc: boolean vector of size 3 storing if pbc is enabled for that direction.
            cutoff: the cutoff inside which atoms are considered as pairs

        Returns:
            numbers of repeats required to make the repeated box contains all the neighbors
            of atoms in the ceter cell
        """
        reciprocal_cell = cell.inverse().t()
        inv_distances = reciprocal_cell.norm(2, -1)
        result = torch.ceil(cutoff * inv_distances).to(torch.long)
        return torch.where(pbc, result, torch.tensor(0))


    def map2central(cell: Tensor, coordinates: Tensor, pbc: Tensor) -> Tensor:
        """Map atoms outside the unit cell into the cell using PBC.

        Arguments:
            cell: tensor of shape ``(3, 3)`` of the three vectors defining unit cell:

                .. code-block:: python

                    tensor([[x1, y1, z1],
                            [x2, y2, z2],
                            [x3, y3, z3]])

            coordinates: Tensor of shape ``(atoms, 3)`` or ``(molecules, atoms, 3)``.
            pbc: boolean vector of size 3 storing if pbc is enabled for that direction.

        Returns:
            coordinates of atoms mapped back to unit cell.
        """
        # Step 1: convert coordinates from standard cartesian coordinate to unit
        # cell coordinates
        inv_cell = torch.inverse(cell)
        coordinates_cell = coordinates @ inv_cell
        # Step 2: wrap cell coordinates into [0, 1)
        coordinates_cell -= coordinates_cell.floor() * pbc.to(coordinates_cell.dtype)
        # Step 3: convert from cell coordinates back to standard cartesian
        # coordinate
        return coordinates_cell @ cell








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

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


.. _sphx_glr_download_code_pbc.py:


.. only :: html

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



  .. container:: sphx-glr-download

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



  .. container:: sphx-glr-download

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


.. only:: html

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

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