.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/1-charges-example.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_1-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-40 .. code-block:: Python import torch import vesin.torch from metatensor.torch import Labels, TensorBlock from metatensor.torch.atomistic import System import torchpme .. GENERATED FROM PYTHON SOURCE LINES 41-42 Create the properties CsCl unit cell .. GENERATED FROM PYTHON SOURCE LINES 43-51 .. code-block:: Python symbols = ("Cs", "Cl") types = torch.tensor([55, 17]) positions = torch.tensor([(0, 0, 0), (0.5, 0.5, 0.5)], dtype=torch.float64) cell = torch.eye(3, dtype=torch.float64) pbc = torch.tensor([True, True, True]) .. GENERATED FROM PYTHON SOURCE LINES 52-55 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 56-61 .. code-block:: Python smearing, pme_params, cutoff = torchpme.utils.tune_pme( sum_squared_charges=2.0, cell=cell, positions=positions ) .. GENERATED FROM PYTHON SOURCE LINES 62-63 The tuning found the following best values for our system. .. GENERATED FROM PYTHON SOURCE LINES 64-69 .. code-block:: Python print("smearing:", smearing) print("PME parameters:", pme_params) print("cutoff:", cutoff) .. rst-class:: sphx-glr-script-out .. code-block:: none smearing: 0.7967574064355606 PME parameters: {'mesh_spacing': 0.5049053084417655, 'interpolation_nodes': 4} cutoff: 3.849277348506204 .. GENERATED FROM PYTHON SOURCE LINES 70-73 Based on the system we compute the corresponding half neighbor list using `vesin `_ and rearrange the results to be suitable for the calculations below. .. GENERATED FROM PYTHON SOURCE LINES 74-83 .. code-block:: Python nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False) i, j, S, D, neighbor_distances = nl.compute( points=positions, box=cell, periodic=True, quantities="ijSDd" ) neighbor_indices = torch.stack([i, j], dim=1) .. GENERATED FROM PYTHON SOURCE LINES 84-87 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 88-93 .. code-block:: Python calculator = torchpme.PMECalculator( torchpme.CoulombPotential(smearing=smearing), **pme_params ) .. GENERATED FROM PYTHON SOURCE LINES 94-99 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 100-103 .. code-block:: Python charges = torch.tensor([[1.0], [-1.0]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 104-107 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 108-111 .. code-block:: Python print(charges.shape) .. rst-class:: sphx-glr-script-out .. code-block:: none torch.Size([2, 1]) .. GENERATED FROM PYTHON SOURCE LINES 112-113 Calculate the potential using the PMECalculator calculator .. GENERATED FROM PYTHON SOURCE LINES 114-123 .. code-block:: Python potential = calculator( positions=positions, cell=cell, charges=charges, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, ) .. GENERATED FROM PYTHON SOURCE LINES 124-126 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 127-130 .. 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 131-147 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 148-151 .. code-block:: Python charges_one_hot = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 152-160 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 161-170 .. 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 171-173 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 174-177 .. code-block:: Python print(potential_one_hot) .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([[-1.4187, -0.4010], [-0.4010, -1.4187]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 178-180 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 181-186 .. 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 187-191 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 192-197 .. code-block:: Python calculator_metatensor = torchpme.metatensor.PMECalculator( torchpme.CoulombPotential(smearing=smearing), **pme_params ) .. GENERATED FROM PYTHON SOURCE LINES 198-207 Computation with metatensor involves using Metatensor's :class:`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 208-211 .. code-block:: Python system = System(types=types, positions=positions, cell=cell, pbc=pbc) .. GENERATED FROM PYTHON SOURCE LINES 212-214 We first add the neighbor list to the ``system``. This requires creating a ``NeighborList`` object to store the *neighbor indices*, *distances*, and *shifts*. .. GENERATED FROM PYTHON SOURCE LINES 215-232 .. code-block:: Python sample_values = torch.hstack([neighbor_indices, S]) samples = Labels( names=[ "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c", ], values=sample_values, ) components = Labels(names=["xyz"], values=torch.tensor([[0, 1, 2]]).T) properties = Labels(names=["distance"], values=torch.tensor([[0]])) neighbors = TensorBlock(D.reshape(-1, 3, 1), samples, [components], properties) .. GENERATED FROM PYTHON SOURCE LINES 233-240 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 `. .. GENERATED FROM PYTHON SOURCE LINES 241-250 .. code-block:: Python # Create a TensorBlock for the charges data = TensorBlock( values=charges, samples=Labels.range("atom", charges.shape[0]), components=[], properties=Labels.range("charge", charges.shape[1]), ) .. GENERATED FROM PYTHON SOURCE LINES 251-252 Add the charges data to the system .. GENERATED FROM PYTHON SOURCE LINES 253-256 .. code-block:: Python system.add_data(name="charges", data=data) .. GENERATED FROM PYTHON SOURCE LINES 257-258 We now calculate the potential using the MetaPMECalculator calculator .. GENERATED FROM PYTHON SOURCE LINES 259-262 .. 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:101: 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:865.) n_charges_channels = system.get_data("charges").values.shape[1] .. GENERATED FROM PYTHON SOURCE LINES 263-265 The calculated potential is wrapped inside a :class:`TensorMap ` and annotated with metadata of the computation. .. GENERATED FROM PYTHON SOURCE LINES 266-269 .. 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 270-272 The tensorMap has *1* :class:`TensorBlock ` and the values of the potential are stored in the ``values`` property. .. GENERATED FROM PYTHON SOURCE LINES 273-276 .. code-block:: Python print(potential_metatensor[0].values) .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([[-1.0177], [ 1.0177]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 277-280 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 281-284 .. 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 285-296 If you want to inspect the metadata in more detail, you can access the :class:`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 ` method. .. GENERATED FROM PYTHON SOURCE LINES 297-308 .. code-block:: Python data_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]), ) # Add the charges data to the system system.add_data(name="charges", data=data_one_hot, override=True) .. GENERATED FROM PYTHON SOURCE LINES 309-310 Finally, we calculate the potential using ``calculator_metatensor`` .. GENERATED FROM PYTHON SOURCE LINES 311-313 .. code-block:: Python potential = calculator_metatensor.forward(system, neighbors) .. GENERATED FROM PYTHON SOURCE LINES 314-315 And as above, the values of the potential are the same. .. GENERATED FROM PYTHON SOURCE LINES 316-319 .. code-block:: Python print(potential[0].values) .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([[-1.4187, -0.4010], [-0.4010, -1.4187]], dtype=torch.float64) .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 12.898 seconds) .. _sphx_glr_download_examples_1-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: 1-charges-example.ipynb <1-charges-example.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 1-charges-example.py <1-charges-example.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 1-charges-example.zip <1-charges-example.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_