.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/09-atomistic-model.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_09-atomistic-model.py: Atomistic model for molecular dynamics ====================================== .. currentmodule:: torchpme In this example, we demonstrate how to construct a `metatensor atomistic model `_ based on the :ref:`metatensor interface ` of ``torchpme``. The model will be used to run a very short molecular dynamics (MD) simulation of a non-neutral hydroden plasma in a cubic box. The plasma consists of massive point particles which are interacting pairwise via a Coulomb force which we compute using the :class:`EwaldCalculator`. The tutorial assumes knowledge of ``torchpme`` and how to export an metatensor atomistic model to run it with the ASE calculator. For learning these details we refer to the `metatensor atomistic tutorials `_. .. GENERATED FROM PYTHON SOURCE LINES 21-52 .. code-block:: Python from typing import Dict, List, Optional # noqa # tools to run the simulation and visualization import ase.md import ase.visualize.plot import ase.md.velocitydistribution import chemiscope import metatensor.torch import matplotlib.pyplot as plt # tools to wrap and run the model import numpy as np import torch from metatensor.torch import Labels, TensorBlock, TensorMap from metatensor.torch.atomistic import ( MetatensorAtomisticModel, ModelCapabilities, ModelEvaluationOptions, ModelMetadata, ModelOutput, NeighborListOptions, System, ) # Integration with ASE calculator for metatensor atomistic models from metatensor.torch.atomistic.ase_calculator import MetatensorCalculator # the usual suspect import torchpme .. GENERATED FROM PYTHON SOURCE LINES 53-58 The simulation system --------------------- Create a system of 12 hydrogen atoms in a cubic periodic box of :math:`10\,\text{Å}` side length. .. GENERATED FROM PYTHON SOURCE LINES 59-69 .. code-block:: Python rng = np.random.default_rng(42) atoms = ase.Atoms( 12 * "H", positions=10 * rng.random([12, 3]), cell=10 * np.eye(3), pbc=True, ) .. GENERATED FROM PYTHON SOURCE LINES 70-71 We now can visualize the system with `chemiscope `_. .. GENERATED FROM PYTHON SOURCE LINES 72-81 .. code-block:: Python chemiscope.show( [atoms], mode="structure", settings=chemiscope.quick_settings( trajectory=True, structure_settings={"unitCell": True} ), ) .. chemiscope:: _datasets/fig_09-atomistic-model_010.json.gz :mode: structure .. raw:: html


.. GENERATED FROM PYTHON SOURCE LINES 82-91 Assigning charges ----------------- For computing the electrostatic potential we need to assign ``charges`` to a :class:`metatensor.torch.atomistic.System` since ``charges`` will not be provided by the engine. Here, we define a simple charge assignment scheme based on the atomic number. We set the partial charge of all hydrogens :math:`1` in terms of the elementary charge. Such an assignemnt scheme can also be more complex and for example deduce the charges from a short range machine learning model. .. GENERATED FROM PYTHON SOURCE LINES 92-115 .. code-block:: Python def add_charges(system: System) -> None: dtype = system.positions.dtype device = system.positions.device # set charges of all atoms to 1 charges = torch.ones(len(system), 1, dtype=dtype, device=device) # Create metadata for the charges TensorBlock samples = Labels("atom", torch.arange(len(system), device=device).reshape(-1, 1)) properties = Labels("charge", torch.zeros(1, 1, device=device, dtype=torch.int32)) block = TensorBlock( values=charges, samples=samples, components=[], properties=properties ) tensor = TensorMap( keys=Labels("_", torch.zeros(1, 1, dtype=torch.int32)), blocks=[block] ) system.add_data(name="charges", tensor=tensor) .. GENERATED FROM PYTHON SOURCE LINES 116-128 .. warning:: Having non-charge neutral system is not problematic for this simulation. Even though any Ewald method requires a charge-neutral system, charged system will be neutralized by adding an homogenous background density. This background charge adds an additional contribution to the real-space interaction which is already accountaed for by torchpme. For homogenous and isotropic systems a nonzero background charge will not have an effect on the simulation but it will for inhomgenous system. For more information see e.g. `Hub et.al `_. We now test the assignmenet by creating a test ``system`` and adding the charges. .. GENERATED FROM PYTHON SOURCE LINES 129-141 .. code-block:: Python system = System( types=torch.from_numpy(atoms.get_atomic_numbers()), positions=torch.from_numpy(atoms.positions), cell=torch.from_numpy(atoms.cell.array), pbc=torch.from_numpy(atoms.pbc), ) add_charges(system) print(system.get_data("charges").block().values) .. rst-class:: sphx-glr-script-out .. code-block:: none tensor([[1.], [1.], [1.], [1.], [1.], [1.], [1.], [1.], [1.], [1.], [1.], [1.]], dtype=torch.float64) .. GENERATED FROM PYTHON SOURCE LINES 142-155 As expected the charges are assigned to all atoms based on their atomic number. The model --------- We now define the atomistic model that computes the potential based on a given ``calculator`` and the ``cutoff``. The ``cutoff`` is required to define the neighbor list which will be used to compute the short range part of the potential. The design of the model is inspire by the `Lennard-Jones model `_. Inside the ``forward`` method we compute the potential at the atomci positions, which is then multiplied by the formal "charges" to obtain a per-atom energy. .. GENERATED FROM PYTHON SOURCE LINES 156-248 .. code-block:: Python class CalculatorModel(torch.nn.Module): def __init__(self, calculator: torchpme.metatensor.Calculator, cutoff: float): super().__init__() self.calculator = calculator # We use a half neighborlist and allow to have pairs farther than cutoff # (`strict=False`) since this is not problematic for PME and may speed up the # computation of the neigbors. self.nl = NeighborListOptions(cutoff=cutoff, full_list=False, strict=False) def requested_neighbor_lists(self): return [self.nl] def _setup_systems( self, systems: list[System], selected_atoms: Optional[Labels] = None, ) -> tuple[System, TensorBlock]: """Remove possible ghost atoms and add charges to the system.""" if len(systems) > 1: raise ValueError(f"only one system supported, got {len(systems)}") system_i = 0 system = systems[system_i] # select only real atoms and discard ghosts if selected_atoms is not None: current_system_mask = selected_atoms.column("system") == system_i current_atoms = selected_atoms.column("atom") current_atoms = current_atoms[current_system_mask].to(torch.long) types = system.types[current_atoms] positions = system.positions[current_atoms] else: types = system.types positions = system.positions system_final = System(types, positions, system.cell, system.pbc) add_charges(system_final) return system_final, system.get_neighbor_list(self.nl) def forward( self, systems: List[System], # noqa outputs: Dict[str, ModelOutput], # noqa selected_atoms: Optional[Labels] = None, ) -> Dict[str, TensorMap]: # noqa if list(outputs.keys()) != ["energy"]: raise ValueError( f"`outputs` keys ({', '.join(outputs.keys())}) contain unsupported " "keys. Only 'energy' is supported." ) system, neighbors = self._setup_systems(systems, selected_atoms) # compute the potential using torchpme potential = self.calculator.forward(system, neighbors) # Create a reference charge block with the same metadata as the potential to # allow multiplcation which requries same metadata charges_block = TensorBlock( values=system.get_data("charges").block().values, samples=potential[0].samples, components=potential[0].components, properties=potential[0].properties, ) charge_map = TensorMap(keys=potential.keys, blocks=[charges_block]) energy_per_atom = metatensor.torch.multiply(potential, charge_map) if outputs["energy"].per_atom: energy = energy_per_atom else: energy = metatensor.torch.sum_over_samples( energy_per_atom, sample_names="atom" ) # Rename property label to follow metatensor's covention for an atomistic model old_block = energy[0] block = TensorBlock( values=old_block.values, samples=old_block.samples, components=old_block.components, properties=old_block.properties.rename("charges_channel", "energy"), ) return {"energy": TensorMap(keys=energy.keys, blocks=[block])} .. GENERATED FROM PYTHON SOURCE LINES 249-273 .. warning:: Due to limitatations in the engine interface of the :class:`MetatensorAtomisticModel `, the evaluation of the energy for a subset of atoms is not supported. If you want to compute the energy for a subset you have to filter the contributions after the computation of the whole system. Define a calculator ------------------- To test the model we need to define a calculator that computes the potential. We here use a the :class:`metatensor.EwaldCalculator` and a :class:`CoulombPotential`. .. note:: We here use the Ewald method and PME because the system only contains 16 atoms. For such small systems the Ewald method is up to 6 times faster compared to PME. If your system reaches a size of 1000 atoms it is recommended to use :class:`metatensor.PMECalculator`, or :class:`metatensor.P3MCalculator` that implements the particle-particle/particle-mesh method. See at the end of this tutorial for an example. These are rather tight settings you can try :func:`tune_ewald ` to determine automatically parameters with a target accuracy .. GENERATED FROM PYTHON SOURCE LINES 274-277 .. code-block:: Python smearing, ewald_params, cutoff = 8.0, {"lr_wavelength": 64.0}, 32.0 .. GENERATED FROM PYTHON SOURCE LINES 278-279 We now define an Ewald calculator with a Coulomb potential. .. GENERATED FROM PYTHON SOURCE LINES 280-287 .. code-block:: Python calculator = torchpme.metatensor.EwaldCalculator( torchpme.CoulombPotential(smearing=smearing), **ewald_params, prefactor=torchpme.prefactors.eV_A, ) .. GENERATED FROM PYTHON SOURCE LINES 288-295 Define model metatdata ---------------------- We now initilize the model and wrap it in a metatensor atomistic model defining all necessary metatdata. This contains the (energy) units and length units. .. GENERATED FROM PYTHON SOURCE LINES 296-300 .. code-block:: Python energy_unit = "eV" length_unit = "angstrom" .. GENERATED FROM PYTHON SOURCE LINES 301-307 We now have to define what our model is able to compute. The :class:`metatensor.torch.atomistic.ModelOutput` defines that the model will compute the energy in eV. Besides the obvous parameters (``atomic_types``, ``supported_devices`` and the ``dtype``) it is very important to set ``interaction_range`` to be **infinite**. For a finite range the provided system of the engine might only contain a subset of the system which will lead to wrong results! .. GENERATED FROM PYTHON SOURCE LINES 308-323 .. code-block:: Python outputs = {"energy": ModelOutput(quantity="energy", unit=energy_unit, per_atom=False)} options = ModelEvaluationOptions( length_unit=length_unit, outputs=outputs, selected_atoms=None ) model_capabilities = ModelCapabilities( outputs=outputs, atomic_types=[1], interaction_range=torch.inf, length_unit=length_unit, supported_devices=["cpu", "cuda"], dtype="float32", ) .. GENERATED FROM PYTHON SOURCE LINES 324-326 Initilize and wrap the model ---------------------------- .. GENERATED FROM PYTHON SOURCE LINES 327-336 .. code-block:: Python model = CalculatorModel(calculator=calculator, cutoff=cutoff) model.eval() atomistic_model = MetatensorAtomisticModel( model.eval(), ModelMetadata(), model_capabilities ) .. GENERATED FROM PYTHON SOURCE LINES 337-344 We'll run the simulation in the constant volume/temperature thermodynamic ensemble (NVT or Canonical ensemble), using a Langevin thermostat for time integration. Please refer to the corresponding documentation (:class:`ase.md.langevin.Langevin`) for more information! To start the simulation we first set the ``atomistic_model`` as the calculator for our plasma. .. GENERATED FROM PYTHON SOURCE LINES 345-349 .. code-block:: Python ewald_mta_calculator = MetatensorCalculator(atomistic_model) atoms.calc = ewald_mta_calculator .. GENERATED FROM PYTHON SOURCE LINES 350-351 Set initial velocities according to the Maxwell-Boltzmann distribution .. GENERATED FROM PYTHON SOURCE LINES 352-355 .. code-block:: Python ase.md.velocitydistribution.MaxwellBoltzmannDistribution(atoms, 10_000 * ase.units.kB) .. rst-class:: sphx-glr-script-out .. code-block:: none /home/runner/work/torch-pme/torch-pme/.tox/docs/lib/python3.13/site-packages/ase/md/md.py:52: FutureWarning: Specify the temperature in K using the 'temperature_K' argument warnings.warn(FutureWarning(w)) .. GENERATED FROM PYTHON SOURCE LINES 356-357 Set up the Langevin thermostat for NVT ensemble .. GENERATED FROM PYTHON SOURCE LINES 358-366 .. code-block:: Python integrator = ase.md.Langevin( atoms, timestep=2 * ase.units.fs, temperature_K=10_000, friction=0.1 / ase.units.fs, ) .. GENERATED FROM PYTHON SOURCE LINES 367-373 Run the simulation ------------------ We now have everything in place run the simulation for 50 steps (:math:`2\,\mathrm{fs}`) and collect the potential, kinetic and total energy as well as the temperature and pressure. .. GENERATED FROM PYTHON SOURCE LINES 374-395 .. code-block:: Python n_steps = 500 potential_energy = np.zeros(n_steps) kinetic_energy = np.zeros(n_steps) total_energy = np.zeros(n_steps) temperature = np.zeros(n_steps) pressure = np.zeros(n_steps) trajectory = [] for i_step in range(n_steps): integrator.run(1) # collect data about the simulation trajectory.append(atoms.copy()) potential_energy[i_step] = atoms.get_potential_energy() kinetic_energy[i_step] = atoms.get_kinetic_energy() total_energy[i_step] = atoms.get_total_energy() temperature[i_step] = atoms.get_temperature() pressure[i_step] = -np.diagonal(atoms.get_stress(voigt=False)).mean() .. GENERATED FROM PYTHON SOURCE LINES 396-398 We can now use `chemiscope `_ to visualize the trajectory. For better visualization we wrap the atoms inside the unit cell. .. GENERATED FROM PYTHON SOURCE LINES 399-411 .. code-block:: Python for atoms in trajectory: atoms.wrap() chemiscope.show( trajectory, mode="structure", settings=chemiscope.quick_settings( trajectory=True, structure_settings={"unitCell": True} ), ) .. chemiscope:: _datasets/fig_09-atomistic-model_011.json.gz :mode: structure .. raw:: html


.. GENERATED FROM PYTHON SOURCE LINES 412-416 Analyze the results ------------------- And look at the time evolution of some physical constants for our system .. GENERATED FROM PYTHON SOURCE LINES 416-440 .. code-block:: Python fig, ax = plt.subplots(3, figsize=(8, 5), sharex=True) time = 2.0 * np.arange(n_steps) ax[0].plot(time, potential_energy, label="potential energy") ax[0].plot(time, kinetic_energy, label="kinetic energy") ax[0].plot(time, total_energy, label="total energy") ax[0].legend(ncol=3) ax[0].set_ylabel("energy [eV]") ax[1].plot(time, temperature, label="temperature") ax[1].axhline(10_000, color="black", linestyle="--", label="target temperature") ax[1].legend(ncol=2) ax[1].set_ylabel("temperature [K]") ax[2].plot(time, pressure) ax[2].set_ylabel("pressure [eV Å$^{-3}$]") ax[-1].set_xlabel("time / fs") fig.align_labels() plt.show() .. image-sg:: /examples/images/sphx_glr_09-atomistic-model_001.png :alt: 09 atomistic model :srcset: /examples/images/sphx_glr_09-atomistic-model_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 441-454 Given the presence of a Langevin thermostat, the total energy is not conserved, but the temperature is well-controlled and fluctuates around the target value of 10,000 K. The metatensor interface also is able to compute the pressure of the simulation via auto differentiation, which is plotted as well. If you want to know more about thermostats and constant-temperature molecular dynamics, you can see `this tutorial `_. This atomistic model can also be used in other engines like LAMMPS. See the metatensor atomistic page on `supported simulation engines `_. The presented model can also be used in more complex sitatuations, e.g. to compute only the electrostatic potential while other parts of the simulations such as the Lennard-Jones potential are computed by other calculators inside the simulation engine. .. GENERATED FROM PYTHON SOURCE LINES 457-463 A comparison between calculators ================================ Even though, as discussed above, for such a small simulation the Ewald method is likely to be the most efficient, it is easy to set up a model that uses the PME, or P3M, calculators in ``torchpme``. For example, .. GENERATED FROM PYTHON SOURCE LINES 463-500 .. code-block:: Python # PME smearing, ewald_params, cutoff = ( 0.5, {"mesh_spacing": 0.25, "interpolation_nodes": 4}, 8.0, ) pme_calculator = torchpme.metatensor.PMECalculator( torchpme.CoulombPotential(smearing=smearing), **ewald_params, prefactor=torchpme.prefactors.eV_A, ) pme_model = CalculatorModel(calculator=pme_calculator, cutoff=cutoff) pme_atomistic_model = MetatensorAtomisticModel( pme_model.eval(), ModelMetadata(), model_capabilities ) pme_mta_calculator = MetatensorCalculator(pme_atomistic_model) # P3M p3m_calculator = torchpme.metatensor.P3MCalculator( torchpme.CoulombPotential(smearing=smearing), **ewald_params, prefactor=torchpme.prefactors.eV_A, ) p3m_model = CalculatorModel(calculator=p3m_calculator, cutoff=cutoff) p3m_atomistic_model = MetatensorAtomisticModel( p3m_model.eval(), ModelMetadata(), model_capabilities ) p3m_mta_calculator = MetatensorCalculator(p3m_atomistic_model) .. GENERATED FROM PYTHON SOURCE LINES 501-502 We can then compare the values obtained with the two Ewald engines .. GENERATED FROM PYTHON SOURCE LINES 502-532 .. code-block:: Python # gets the energy from the Ewald calculator atoms.calc = ewald_mta_calculator ewald_energy = atoms.get_potential_energy() ewald_forces = atoms.get_forces() # overrides the calculator and computes PME atoms = atoms.copy() atoms.calc = pme_mta_calculator pme_energy = atoms.get_potential_energy() pme_forces = atoms.get_forces() # ... and P3M atoms = atoms.copy() atoms.calc = p3m_mta_calculator p3m_energy = atoms.get_potential_energy() p3m_forces = atoms.get_forces() print( f"Energy (Ewald): {ewald_energy}\n" f"Energy (PME): {pme_energy}\n" f"Energy (P3M): {p3m_energy}\n" ) print( f"Forces (Ewald):\n{ewald_forces}\n\n" f"Forces (PME):\n{pme_forces}\n\n" f"Forces (P3M):\n{p3m_forces}" ) .. rst-class:: sphx-glr-script-out .. code-block:: none Energy (Ewald): -50.571353912353516 Energy (PME): -50.52611541748047 Energy (P3M): -50.52189254760742 Forces (Ewald): [[ 0.09059066 0.44978419 -0.90444809] [-0.37997434 -2.0111928 1.07122648] [ 0.18100952 -0.43447021 -0.03110215] [ 1.462111 1.1936059 0.03581454] [-0.85438287 -0.05459023 0.40494722] [-0.17709282 -0.473171 -0.11635519] [-0.41979057 0.39633575 -0.47465798] [-0.05026046 -0.53184015 -0.22202076] [ 0.34676024 0.19488281 0.14477026] [ 0.74841189 0.34315687 0.49568146] [ 1.20114195 1.10904992 -0.30248025] [-2.14852881 -0.18155357 -0.10137586]] Forces (PME): [[ 0.09481248 0.44857794 -0.90494817] [-0.38760597 -2.00512981 1.07789755] [ 0.1772218 -0.44207641 -0.03918238] [ 1.45665932 1.18973684 0.03450844] [-0.84592515 -0.05447095 0.40587896] [-0.17136928 -0.46527734 -0.11410771] [-0.42817628 0.38999006 -0.4783752 ] [-0.0519303 -0.52337688 -0.21536151] [ 0.3421914 0.2006797 0.1378473 ] [ 0.75634819 0.34439543 0.49703431] [ 1.20872116 1.10101509 -0.30976978] [-2.14461827 -0.1839259 -0.09277071]] Forces (P3M): [[ 0.09129928 0.44955575 -0.90452892] [-0.38069603 -2.01033592 1.07196772] [ 0.18033987 -0.43509534 -0.03153744] [ 1.46141398 1.19291496 0.03553772] [-0.85430956 -0.05454677 0.40515703] [-0.17627423 -0.4728452 -0.11593961] [-0.42024806 0.39555588 -0.47535861] [-0.05051793 -0.53170562 -0.22120857] [ 0.34617156 0.19568491 0.14406092] [ 0.74886841 0.3434746 0.4960297 ] [ 1.2017554 1.1085912 -0.30324963] [-2.14785647 -0.181995 -0.10121254]] .. GENERATED FROM PYTHON SOURCE LINES 533-534 The above output shows us that these values are very close to each other, as expected. .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 35.323 seconds) .. _sphx_glr_download_examples_09-atomistic-model.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 09-atomistic-model.ipynb <09-atomistic-model.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 09-atomistic-model.py <09-atomistic-model.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 09-atomistic-model.zip <09-atomistic-model.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_