.. DO NOT EDIT.
.. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY.
.. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE:
.. "examples/11-4-site-water.py"
.. LINE NUMBERS ARE GIVEN BELOW.

.. only:: html

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

        :ref:`Go to the end <sphx_glr_download_examples_11-4-site-water.py>`
        to download the full example code.

.. rst-class:: sphx-glr-example-title

.. _sphx_glr_examples_11-4-site-water.py:


.. _example-tip4p-water:

4-site water models
===================

.. currentmodule:: torchpme

# Several water models (starting from the venerable TIP4P model of
# `Abascal and C. Vega, JCP (2005) <http://doi.org/10.1063/1.2121687>`_)
# use a center of negative charge that is displaced from the O position.
# This is easily implemented, yielding the forces on the O and H positions
# generated by the displaced charge.

.. GENERATED FROM PYTHON SOURCE LINES 15-34

.. code-block:: Python


    import ase
    import torch

    import torchpme

    structure = ase.Atoms(
        positions=[
            [0, 0, 0],
            [0, 1, 0],
            [1, -0.2, 0],
        ],
        cell=[6, 6, 6],
        symbols="OHH",
    )

    cell = torch.from_numpy(structure.cell.array)
    positions = torch.from_numpy(structure.positions)








.. GENERATED FROM PYTHON SOURCE LINES 35-37

The key step is to create a "fourth site" based on the oxygen positions and use it in
the ``interpolate`` step.

.. GENERATED FROM PYTHON SOURCE LINES 37-52

.. code-block:: Python


    charges = torch.tensor([[-1.0], [0.5], [0.5]])

    positions.requires_grad_(True)
    charges.requires_grad_(True)
    cell.requires_grad_(True)

    positions_4site = torch.vstack(
        [
            ((positions[1::3] + positions[2::3]) * 0.5 + positions[0::3] * 3) / 4,
            positions[1::3],
            positions[2::3],
        ]
    )








.. GENERATED FROM PYTHON SOURCE LINES 53-57

.. important::

  For the automatic differentiation to work it is important to make a new tensor as
  ``positions_4site`` and do not "overwrite" the original tensor.

.. GENERATED FROM PYTHON SOURCE LINES 57-67

.. code-block:: Python


    ns = torch.tensor([5, 5, 5])
    interpolator = torchpme.lib.MeshInterpolator(
        cell=cell, ns_mesh=ns, interpolation_nodes=3, method="Lagrange"
    )
    interpolator.compute_weights(positions_4site)
    mesh = interpolator.points_to_mesh(charges)

    value = (mesh**2).sum()








.. GENERATED FROM PYTHON SOURCE LINES 68-70

The gradients can be computed by just running `backward` on the
end result. Gradients are computed on the H and O positions.

.. GENERATED FROM PYTHON SOURCE LINES 70-85

.. code-block:: Python


    value.backward()

    print(
        f"""
    Position gradients:
    {positions.grad.T}

    Cell gradients:
    {cell.grad}

    Charges gradients:
    {charges.grad.T}
    """
    )




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

 .. code-block:: none


    Position gradients:
    tensor([[-0.5106, -0.0857,  0.5133],
            [-0.4691,  0.5263, -0.0164],
            [ 0.0000,  0.0000,  0.0000]], dtype=torch.float64)

    Cell gradients:
    tensor([[-0.0855,  0.0027, -0.0000],
            [ 0.0314, -0.0883, -0.0000],
            [-0.0000, -0.0000, -0.0000]], dtype=torch.float64)

    Charges gradients:
    tensor([[-1.6628,  0.6843,  0.6317]])






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

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


.. _sphx_glr_download_examples_11-4-site-water.py:

.. only:: html

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

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

      :download:`Download Jupyter notebook: 11-4-site-water.ipynb <11-4-site-water.ipynb>`

    .. container:: sphx-glr-download sphx-glr-download-python

      :download:`Download Python source code: 11-4-site-water.py <11-4-site-water.py>`

    .. container:: sphx-glr-download sphx-glr-download-zip

      :download:`Download zipped: 11-4-site-water.zip <11-4-site-water.zip>`


.. only:: html

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

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