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

.. only:: html

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

        :ref:`Go to the end <sphx_glr_download_examples_01-charges-example.py>`
        to download the full example code.

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

.. _sphx_glr_examples_01-charges-example.py:


Computations with Multiple Charge Channels
==========================================

.. currentmodule:: torchpme

In a physical system, the (electrical) charge is a scalar atomic property, and besides
the distance between the particles, the charge defines the electrostatic potential. When
computing a potential with Meshlode, you can not only pass a (reshaped) 1-D array
containing the charges to the ``compute`` method of calculators, but you can also pass a
2-D array containing multiple charges per atom. Meshlode will then calculate one
potential per so-called *charge-channel*. For the standard electrostatic potential, the
number of charge channels is *1*. Additional charge channels are especially useful in a
machine learning task where, for example, one wants to create one potential for each
species in an atomistic system using a so-called *one-hot encoding* of charges.

Here, we will demonstrate how to use the ability of multiple charge channels for a
*CsCl* crystal, where we will cover both the Torch and Metatensor interfaces of
Meshlode.

Torch Version
--------------
First, we will work with the Torch version of Meshlode. This involves using `PyTorch`_
for tensor operations and `ASE`_ (Atomic Simulation Environment) for creating and
manipulating atomic structures.

.. _`PyTorch`: https://pytorch.org
.. _`ASE`: https://wiki.fysik.dtu.dk/ase

.. GENERATED FROM PYTHON SOURCE LINES 32-44

.. code-block:: Python


    import torch
    import vesin.torch
    import vesin.torch.metatensor
    from metatensor.torch import Labels, TensorBlock, TensorMap
    from metatensor.torch.atomistic import NeighborListOptions, System

    import torchpme
    from torchpme.tuning import tune_pme

    dtype = torch.float64








.. GENERATED FROM PYTHON SOURCE LINES 45-46

Create the properties CsCl unit cell

.. GENERATED FROM PYTHON SOURCE LINES 47-55

.. code-block:: Python

    symbols = ("Cs", "Cl")
    types = torch.tensor([55, 17])
    charges = torch.tensor([[1.0], [-1.0]], dtype=dtype)
    positions = torch.tensor([(0, 0, 0), (0.5, 0.5, 0.5)], dtype=dtype)
    cell = torch.eye(3, dtype=dtype)
    pbc = torch.tensor([True, True, True])









.. GENERATED FROM PYTHON SOURCE LINES 56-59

Based on our system we will first *tune* the PME parameters for an accurate computation.
The ``sum_squared_charges`` is equal to ``2.0`` becaue each atom either has a charge
of 1 or -1 in units of elementary charges.

.. GENERATED FROM PYTHON SOURCE LINES 60-78

.. code-block:: Python


    cutoff = 4.4
    nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False)
    neighbor_indices, neighbor_distances = nl.compute(
        points=positions.to(dtype=torch.float64, device="cpu"),
        box=cell.to(dtype=torch.float64, device="cpu"),
        periodic=True,
        quantities="Pd",
    )
    smearing, pme_params, _ = tune_pme(
        charges=charges,
        cell=cell,
        positions=positions,
        cutoff=cutoff,
        neighbor_indices=neighbor_indices,
        neighbor_distances=neighbor_distances,
    )








.. GENERATED FROM PYTHON SOURCE LINES 79-80

The tuning found the following best values for our system.

.. GENERATED FROM PYTHON SOURCE LINES 81-86

.. code-block:: Python


    print("smearing:", smearing)
    print("PME parameters:", pme_params)
    print("cutoff:", cutoff)





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

 .. code-block:: none

    smearing: 1.1069526756106463
    PME parameters: {'interpolation_nodes': 3, 'mesh_spacing': 0.2857142857142857}
    cutoff: 4.4




.. GENERATED FROM PYTHON SOURCE LINES 87-90

Based on the system we compute the corresponding half neighbor list using `vesin
<https://luthaf.fr/vesin>`_ and rearrange the results to be suitable for the
calculations below.

.. GENERATED FROM PYTHON SOURCE LINES 91-98

.. code-block:: Python


    nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False)

    neighbor_indices, S, D, neighbor_distances = nl.compute(
        points=positions, box=cell, periodic=True, quantities="PSDd"
    )








.. GENERATED FROM PYTHON SOURCE LINES 99-102

Next, we initialize the :class:`PMECalculator` calculator with an ``exponent`` of
*1* for electrostatic interactions between the two atoms. This calculator
will be used to *compute* the potential energy of the system.

.. GENERATED FROM PYTHON SOURCE LINES 103-108

.. code-block:: Python


    calculator = torchpme.PMECalculator(
        torchpme.CoulombPotential(smearing=smearing), **pme_params
    )
    calculator.to(dtype=dtype)




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

 .. code-block:: none


    PMECalculator(
      (potential): CoulombPotential()
      (kspace_filter): KSpaceFilter(
        (kernel): CoulombPotential()
      )
      (mesh_interpolator): MeshInterpolator()
    )



.. GENERATED FROM PYTHON SOURCE LINES 109-114

Single Charge Channel
#####################

As a first application of multiple charge channels, we start simply by using the
classic definition of one charge channel per atom.

.. GENERATED FROM PYTHON SOURCE LINES 115-118

.. code-block:: Python


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








.. GENERATED FROM PYTHON SOURCE LINES 119-122

Any input the calculators has to be a 2D array where the *rows* describe the number of
atoms (here ``(2)``) and the *columns* the number of atomic charge channels (here
``(1)``).

.. GENERATED FROM PYTHON SOURCE LINES 123-126

.. code-block:: Python


    print(charges.shape)





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

 .. code-block:: none

    torch.Size([2, 1])




.. GENERATED FROM PYTHON SOURCE LINES 127-128

Calculate the potential using the PMECalculator calculator

.. GENERATED FROM PYTHON SOURCE LINES 129-138

.. code-block:: Python


    potential = calculator(
        positions=positions,
        cell=cell,
        charges=charges,
        neighbor_indices=neighbor_indices,
        neighbor_distances=neighbor_distances,
    )








.. GENERATED FROM PYTHON SOURCE LINES 139-141

We find a potential that is close to the Madelung constant of a CsCl crystal which is
:math:`2 \cdot 1.76267 / \sqrt{3} \approx 2.0354`.

.. GENERATED FROM PYTHON SOURCE LINES 142-145

.. code-block:: Python


    print(charges.T @ potential)





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

 .. code-block:: none

    tensor([[-2.0354]], dtype=torch.float64)




.. GENERATED FROM PYTHON SOURCE LINES 146-162

Species-wise One-Hot Encoded Charges
####################################

Now we will compute the potential with multiple channels for the charges. We will use
one channel per species and set the charges to *1* if the atomic ``symbol`` fits the
correct channel. This is called one-hot encoding for each species.

One-hot encoding is a powerful technique in machine learning where categorical data is
converted into a binary vector representation. Each category is represented by a
vector with all elements set to zero except for the index corresponding to that
category, which is set to one. This allows the model to easily distinguish between
different categories without imposing any ordinal relationship between them. In the
context of molecular systems, one-hot encoding can be used to represent different
atomic species as separate charge channels, enabling the calculation of
species-specific potentials and facilitating the learning process for machine learning
algorithms.

.. GENERATED FROM PYTHON SOURCE LINES 163-166

.. code-block:: Python


    charges_one_hot = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=dtype)








.. GENERATED FROM PYTHON SOURCE LINES 167-175

While in ``charges`` there was only one row, ``charges_one_hot`` contains two rows
where the first one corresponds to the Na channel and the second one to the Cl
channel. Consequently, the charge of the Na atom in the Cl channel at index ``(0,1)``
of the ``charges_one_hot`` is zero as well as the ``(1,0)`` which corresponds to the
charge of Cl in the Na channel.

We now again calculate the potential using the same :class:`PMECalculator` calculator
using the ``charges_one_hot`` as input.

.. GENERATED FROM PYTHON SOURCE LINES 176-185

.. code-block:: Python


    potential_one_hot = calculator(
        charges=charges_one_hot,
        cell=cell,
        positions=positions,
        neighbor_indices=neighbor_indices,
        neighbor_distances=neighbor_distances,
    )








.. GENERATED FROM PYTHON SOURCE LINES 186-188

Note that the potential has the same shape as the input charges, but there is a finite
potential on the position of the Cl in the Na channel.

.. GENERATED FROM PYTHON SOURCE LINES 189-192

.. code-block:: Python


    print(potential_one_hot)





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

 .. code-block:: none

    tensor([[-1.4191, -0.4014],
            [-0.4014, -1.4191]], dtype=torch.float64)




.. GENERATED FROM PYTHON SOURCE LINES 193-195

From the potential we can recover the Madelung as above by summing the charge channel
contribution multiplying by the actual partial charge of the atoms.

.. GENERATED FROM PYTHON SOURCE LINES 196-201

.. code-block:: Python


    charge_Na = 1.0
    charge_Cl = -1.0
    print(charge_Na * potential_one_hot[0] + charge_Cl * potential_one_hot[1])





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

 .. code-block:: none

    tensor([-1.0177,  1.0177], dtype=torch.float64)




.. GENERATED FROM PYTHON SOURCE LINES 202-206

Metatensor Version
------------------
Next, we will perform the same exercise with the Metatensor interface. This involves
creating a new calculator with the metatensor interface.

.. GENERATED FROM PYTHON SOURCE LINES 207-212

.. code-block:: Python


    calculator_metatensor = torchpme.metatensor.PMECalculator(
        torchpme.CoulombPotential(smearing=smearing), **pme_params
    )
    calculator_metatensor.to(dtype=dtype)




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

 .. code-block:: none


    PMECalculator(
      (_calculator): PMECalculator(
        (potential): CoulombPotential()
        (kspace_filter): KSpaceFilter(
          (kernel): CoulombPotential()
        )
        (mesh_interpolator): MeshInterpolator()
      )
    )



.. GENERATED FROM PYTHON SOURCE LINES 213-222

Computation with metatensor involves using Metatensor's :class:`System
<metatensor.torch.atomistic.System>` class. The ``System`` stores atomic ``types``,
``positions``, and ``cell`` dimensions.

.. note::
   For our calculations, the parameter ``types`` passed to a ``System`` is redundant;
   it will not be directly used to calculate the potential as the potential depends
   only on the charge of the atom, NOT on the atom's type. However, we still have to
   pass them because it is an obligatory parameter to build the `System` class.

.. GENERATED FROM PYTHON SOURCE LINES 223-226

.. code-block:: Python


    system = System(types=types, positions=positions, cell=cell, pbc=pbc)








.. GENERATED FROM PYTHON SOURCE LINES 227-231

We now compute the neighborlist for our ``system`` using the `vesin metatensor
interface <https://luthaf.fr/vesin/latest/metatensor.html>`_. This requires creating a
:class:`NeighborListOptions <metatensor.torch.atomistic.NeighborListOptions>` to set
the cutoff and the type of list.

.. GENERATED FROM PYTHON SOURCE LINES 232-237

.. code-block:: Python


    options = NeighborListOptions(cutoff=4.0, full_list=True, strict=False)
    nl_mts = vesin.torch.metatensor.NeighborList(options, length_unit="Angstrom")
    neighbors = nl_mts.compute(system)








.. GENERATED FROM PYTHON SOURCE LINES 238-247

Now the ``system`` is ready to be used inside the calculators

Single Charge Channel
#####################

For the metatensor branch, charges of the atoms are defined in a tensor format and
attached to the system as a :class:`TensorBlock <metatensor.torch.TensorBlock>`.

Create a :class:`TensorMap <metetensor.torch.TensorMap>` for the charges

.. GENERATED FROM PYTHON SOURCE LINES 248-260

.. code-block:: Python


    block = TensorBlock(
        values=charges,
        samples=Labels.range("atom", charges.shape[0]),
        components=[],
        properties=Labels.range("charge", charges.shape[1]),
    )

    tensor = TensorMap(
        keys=Labels("_", torch.zeros(1, 1, dtype=torch.int32)), blocks=[block]
    )








.. GENERATED FROM PYTHON SOURCE LINES 261-262

Add the charges data to the system

.. GENERATED FROM PYTHON SOURCE LINES 263-266

.. code-block:: Python


    system.add_data(name="charges", tensor=tensor)








.. GENERATED FROM PYTHON SOURCE LINES 267-268

We now calculate the potential using the MetaPMECalculator calculator

.. GENERATED FROM PYTHON SOURCE LINES 269-272

.. code-block:: Python


    potential_metatensor = calculator_metatensor.forward(system, neighbors)





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

 .. code-block:: none

    /home/runner/work/torch-pme/torch-pme/src/torchpme/metatensor/calculator.py:88: UserWarning: custom data 'charges' is experimental, please contact metatensor's developers to add this data as a member of the `System` class (Triggered internally at /project/metatensor-torch/src/atomistic/system.cpp:866.)
      charge_tensor = system.get_data("charges")




.. GENERATED FROM PYTHON SOURCE LINES 273-275

The calculated potential is wrapped inside a :class:`TensorMap
<metatensor.torch.TensorMap>` and annotated with metadata of the computation.

.. GENERATED FROM PYTHON SOURCE LINES 276-279

.. code-block:: Python


    print(potential_metatensor)





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

 .. code-block:: none

    TensorMap with 1 blocks
    keys: _
          0




.. GENERATED FROM PYTHON SOURCE LINES 280-282

The tensorMap has *1* :class:`TensorBlock <metatensor.torch.TensorBlock>` and the
values of the potential are stored in the ``values`` property.

.. GENERATED FROM PYTHON SOURCE LINES 283-286

.. code-block:: Python


    print(potential_metatensor[0].values)





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

 .. code-block:: none

    tensor([[-1.6768],
            [ 1.6768]], dtype=torch.float64)




.. GENERATED FROM PYTHON SOURCE LINES 287-290

The ``values`` are the same results as for the torch interface shown above
The metadata associated with the ``TensorBlock`` tells us that we have 2 ``samples``
which are our two atoms and 1 property which corresponds to one charge channel.

.. GENERATED FROM PYTHON SOURCE LINES 291-294

.. code-block:: Python


    print(potential_metatensor[0])





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

 .. code-block:: none

    TensorBlock
        samples (2): ['system', 'atom']
        components (): []
        properties (1): ['charges_channel']
        gradients: None





.. GENERATED FROM PYTHON SOURCE LINES 295-306

If you want to inspect the metadata in more detail, you can access the
:class:`Labels <metatensor.torch.Labels>` using the
``potential_metatensor[0].properties`` and ``potential_metatensor[0].samples``
attributes.

Species-wise One-Hot Encoded Charges
####################################

We now create new charges data based on the species-wise ``charges_one_hot`` and
overwrite the ``system``'s charges data using ``override=True`` when applying the
:meth:`add_data <metatensor.torch.atomistic.System.add_data>` method.

.. GENERATED FROM PYTHON SOURCE LINES 307-319

.. code-block:: Python


    block_one_hot = TensorBlock(
        values=charges_one_hot,
        samples=Labels.range("atom", charges_one_hot.shape[0]),
        components=[],
        properties=Labels.range("charge", charges_one_hot.shape[1]),
    )

    tensor_one_hot = TensorMap(
        keys=Labels("_", torch.zeros(1, 1, dtype=torch.int32)), blocks=[block_one_hot]
    )








.. GENERATED FROM PYTHON SOURCE LINES 320-322

Add the charges data to the system. We use the ``override=True`` to overwrite the
already existing charge data from above.

.. GENERATED FROM PYTHON SOURCE LINES 323-326

.. code-block:: Python


    system.add_data(name="charges", tensor=tensor_one_hot, override=True)








.. GENERATED FROM PYTHON SOURCE LINES 327-328

Finally, we calculate the potential using ``calculator_metatensor``

.. GENERATED FROM PYTHON SOURCE LINES 329-331

.. code-block:: Python

    potential = calculator_metatensor.forward(system, neighbors)








.. GENERATED FROM PYTHON SOURCE LINES 332-333

And as above, the values of the potential are the same.

.. GENERATED FROM PYTHON SOURCE LINES 334-337

.. code-block:: Python


    print(potential[0].values)





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

 .. code-block:: none

    tensor([[1.3674, 3.0442],
            [3.0442, 1.3674]], dtype=torch.float64)





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

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


.. _sphx_glr_download_examples_01-charges-example.py:

.. only:: html

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

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

      :download:`Download Jupyter notebook: 01-charges-example.ipynb <01-charges-example.ipynb>`

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

      :download:`Download Python source code: 01-charges-example.py <01-charges-example.py>`

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

      :download:`Download zipped: 01-charges-example.zip <01-charges-example.zip>`


.. only:: html

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

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