.. 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 ` 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-42 .. code-block:: Python import torch import vesin.torch import vesin.torch.metatensor from metatensor.torch import Labels, TensorBlock from metatensor.torch.atomistic import NeighborListOptions, System import torchpme from torchpme.tuning import tune_pme .. GENERATED FROM PYTHON SOURCE LINES 43-44 Create the properties CsCl unit cell .. GENERATED FROM PYTHON SOURCE LINES 45-54 .. code-block:: Python symbols = ("Cs", "Cl") types = torch.tensor([55, 17]) charges = torch.tensor([[1.0], [-1.0]], dtype=torch.float64) 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 55-58 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 59-77 .. 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 78-79 The tuning found the following best values for our system. .. GENERATED FROM PYTHON SOURCE LINES 80-85 .. 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 86-89 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 90-97 .. 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 98-101 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 102-107 .. code-block:: Python calculator = torchpme.PMECalculator( torchpme.CoulombPotential(smearing=smearing), **pme_params ) .. GENERATED FROM PYTHON SOURCE LINES 108-113 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 114-117 .. code-block:: Python charges = torch.tensor([[1.0], [-1.0]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 118-121 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 122-125 .. code-block:: Python print(charges.shape) .. rst-class:: sphx-glr-script-out .. code-block:: none torch.Size([2, 1]) .. GENERATED FROM PYTHON SOURCE LINES 126-127 Calculate the potential using the PMECalculator calculator .. GENERATED FROM PYTHON SOURCE LINES 128-137 .. code-block:: Python potential = calculator( positions=positions, cell=cell, charges=charges, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, ) .. GENERATED FROM PYTHON SOURCE LINES 138-140 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 141-144 .. 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 145-161 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 162-165 .. code-block:: Python charges_one_hot = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 166-174 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 175-184 .. 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 185-187 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 188-191 .. 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 192-194 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 195-200 .. 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 201-205 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 206-211 .. code-block:: Python calculator_metatensor = torchpme.metatensor.PMECalculator( torchpme.CoulombPotential(smearing=smearing), **pme_params ) .. GENERATED FROM PYTHON SOURCE LINES 212-221 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 222-225 .. code-block:: Python system = System(types=types, positions=positions, cell=cell, pbc=pbc) .. GENERATED FROM PYTHON SOURCE LINES 226-230 We now compute the neighborlist for our ``system`` using the `vesin metatensor interface `_. This requires creating a :class:`NeighborListOptions ` to set the cutoff and the type of list. .. GENERATED FROM PYTHON SOURCE LINES 231-236 .. 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 237-244 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 245-254 .. 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 255-256 Add the charges data to the system .. GENERATED FROM PYTHON SOURCE LINES 257-260 .. code-block:: Python system.add_data(name="charges", data=data) .. GENERATED FROM PYTHON SOURCE LINES 261-262 We now calculate the potential using the MetaPMECalculator calculator .. GENERATED FROM PYTHON SOURCE LINES 263-266 .. 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 267-269 The calculated potential is wrapped inside a :class:`TensorMap ` and annotated with metadata of the computation. .. GENERATED FROM PYTHON SOURCE LINES 270-273 .. 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 274-276 The tensorMap has *1* :class:`TensorBlock ` and the values of the potential are stored in the ``values`` property. .. GENERATED FROM PYTHON SOURCE LINES 277-280 .. 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 281-284 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 285-288 .. 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 289-300 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 301-312 .. 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 313-314 Finally, we calculate the potential using ``calculator_metatensor`` .. GENERATED FROM PYTHON SOURCE LINES 315-317 .. code-block:: Python potential = calculator_metatensor.forward(system, neighbors) .. GENERATED FROM PYTHON SOURCE LINES 318-319 And as above, the values of the potential are the same. .. GENERATED FROM PYTHON SOURCE LINES 320-323 .. 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.328 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 `_