.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/dispersion/05_jax_fourier_dftd3_crystal.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_dispersion_05_jax_fourier_dftd3_crystal.py: JAX FourierD3: Periodic Particle-Mesh Dispersion Under ``jax.jit`` ================================================================== FourierD3 evaluates the DFT-D3(BJ) correction on a particle mesh, so the dispersion sum carries no pair cutoff. The one real-space cutoff that remains is the short coordination-number list. The JAX binding is the counterpart of :func:`nvalchemiops.torch.interactions.dispersion.fourier_dftd3`, and the argument list is the same minus ``device``. Eager JAX dispatches each pass separately, and that overhead dominates a correction this cheap. Wrapping the call in :func:`jax.jit` is what makes it worth using, so that is what this example concentrates on. In this example you will learn: - How to build ``FourierD3Parameters`` for the JAX API - How to evaluate the correction on a CsCl crystal - How to ``jax.jit`` the call, and which arguments have to be static to do so For mesh sensitivity, see ``04_fourier_dftd3_crystal.py``; the behaviour is identical here. .. important:: This script is intended as an API demonstration. Do not use this script for performance benchmarking; refer to the `benchmarks` folder instead. .. GENERATED FROM PYTHON SOURCE LINES 44-49 Setup ----- The D3 reference parameters are in atomic units, so every length here is in Bohr. Double precision has to be enabled before the first array is created. .. GENERATED FROM PYTHON SOURCE LINES 49-81 .. code-block:: Python from __future__ import annotations import os from pathlib import Path import numpy as np try: import jax import jax.numpy as jnp except ImportError: print( "This example requires JAX. Install with: pip install 'nvalchemi-toolkit-ops[jax]'" ) raise SystemExit(0) from None jax.config.update("jax_enable_x64", True) from nvalchemiops.jax import neighbors # noqa: E402 from nvalchemiops.jax.interactions.dispersion import ( # noqa: E402 FourierD3Parameters, fourier_dftd3, ) BOHR_TO_ANGSTROM = 0.529177210544 HARTREE_TO_EV = 27.211386245981 if jax.default_backend() == "cpu": print("This example needs a GPU backend to be worth running. Skipping.") raise SystemExit(0) .. rst-class:: sphx-glr-script-out .. code-block:: none Warp UserWarning: Kernel cache artifacts from a previous Warp version were found in '/tmp/warp-cache'. These will be ignored. You can safely delete them. .. GENERATED FROM PYTHON SOURCE LINES 82-88 Reference parameters -------------------- The published Grimme tables, cached by ``utils.py`` on first run. The decomposition of the reference tensor happens once on the host, for the species present, and does not depend on the damping parameters --- one instance serves every functional. .. GENERATED FROM PYTHON SOURCE LINES 88-106 .. code-block:: Python param_file = ( Path(os.path.expanduser("~")) / ".cache" / "nvalchemiops" / "dftd3_parameters.pt" ) # The tables ship as a Torch checkpoint, so reading them needs Torch even on the JAX path. # This is the provenance route ``03_jax_dftd3_molecule.py`` uses too; only the parameters # pass through Torch, never the evaluation. import torch # noqa: E402 if param_file.exists(): torch_tables = torch.load(param_file, weights_only=True) else: from utils import extract_dftd3_parameters, save_dftd3_parameters torch_tables = extract_dftd3_parameters() save_dftd3_parameters(torch_tables) tables = {name: value.numpy() for name, value in torch_tables.items()} .. GENERATED FROM PYTHON SOURCE LINES 107-113 A CsCl crystal -------------- Caesium chloride is simple cubic with Cs at the corner and Cl at the body centre. FourierD3 is periodic by construction: there is no open-boundary path, because the mesh sum is over the infinite lattice. .. GENERATED FROM PYTHON SOURCE LINES 113-146 .. code-block:: Python LATTICE_A = 4.119 / BOHR_TO_ANGSTROM # Bohr REPEATS = 4 basis_fractional = np.array([[0.0, 0.0, 0.0], [0.5, 0.5, 0.5]]) basis_numbers = np.array([55, 17], dtype=np.int32) # Cs, Cl offsets = np.stack(np.meshgrid(*[np.arange(REPEATS)] * 3, indexing="ij"), axis=-1) offsets = offsets.reshape(-1, 3).astype(float) fractional = ((offsets[:, None, :] + basis_fractional[None, :, :]) / REPEATS).reshape( -1, 3 ) cell_np = np.eye(3) * (LATTICE_A * REPEATS) positions = jnp.asarray(fractional @ cell_np) numbers = jnp.asarray(np.tile(basis_numbers, REPEATS**3), dtype=jnp.int32) cell = jnp.asarray(cell_np) n_atoms = int(positions.shape[0]) species = sorted(set(np.asarray(numbers).tolist())) params = FourierD3Parameters.from_tables( tables["rcov"].astype(np.float64), tables["r4r2"].astype(np.float64), tables["c6ab"].astype(np.float64), tables["cn_ref"].astype(np.float64), species=species, ) print(f"atoms : {n_atoms}") print(f"cell edge : {float(cell[0, 0]) * BOHR_TO_ANGSTROM:.2f} Angstrom") print(f"species : {species}") print(f"decomposition rank : {params.rank}") print(f"mesh channels needed : {params.n_species * params.rank}") .. rst-class:: sphx-glr-script-out .. code-block:: none atoms : 128 cell edge : 16.48 Angstrom species : [17, 55] decomposition rank : 3 mesh channels needed : 6 .. GENERATED FROM PYTHON SOURCE LINES 147-153 The coordination-number list ---------------------------- Only the coordination numbers need a neighbour list, so its cutoff is short. It must equal ``cutoff``: the counting function reaches zero exactly there, and a list built to a different radius would silently truncate it. .. GENERATED FROM PYTHON SOURCE LINES 153-167 .. code-block:: Python cutoff = 6.0 / BOHR_TO_ANGSTROM # 6 Angstrom, a typical MLFF cutoff, in Bohr pbc = jnp.array([True, True, True]) neighbor_list, neighbor_ptr, unit_shifts = neighbors.neighbor_list( positions, cutoff=cutoff, cell=cell, pbc=pbc, return_neighbor_list=True, ) print(f"\nneighbour cutoff : {cutoff * BOHR_TO_ANGSTROM:.1f} Angstrom") print(f"directed edges : {neighbor_list.shape[1]}") .. rst-class:: sphx-glr-script-out .. code-block:: none neighbour cutoff : 6.0 Angstrom directed edges : 3328 .. GENERATED FROM PYTHON SOURCE LINES 168-177 Evaluating the correction ------------------------- ``cell`` and ``cutoff`` are both required, and exactly one of ``mesh_dimensions`` and ``mesh_spacing`` must be given --- there is no accuracy-based default. Energy has shape ``(num_systems,)``, forces ``(n_atoms, 3)`` and the virial ``(num_systems, 3, 3)``. Forces and the virial are returned explicitly; ``jax.grad`` of the energy is not the route to them here. .. GENERATED FROM PYTHON SOURCE LINES 177-200 .. code-block:: Python damping = dict(a1=0.4289, a2=4.4407, s8=0.7875) # PBE-D3(BJ) common = dict( fourier_d3_params=params, cell=cell, cutoff=cutoff, mesh_dimensions=(32, 32, 32), neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr, unit_shifts=unit_shifts, ) energy, forces, virial = fourier_dftd3( positions, numbers, **damping, compute_virial=True, **common ) print(f"\nenergy : {float(energy[0]):.8f} Hartree") print(f"per atom : {float(energy[0]) / n_atoms * HARTREE_TO_EV * 1e3:.4f} meV") print( f"max |force| : {float(jnp.abs(forces).max()):.3e} Hartree/Bohr (zero by symmetry)" ) print(f"virial trace : {float(jnp.trace(virial[0])):.6e} Hartree") .. rst-class:: sphx-glr-script-out .. code-block:: none energy : -0.97919696 Hartree per atom : -208.1665 meV max |force| : 1.248e-15 Hartree/Bohr (zero by symmetry) virial trace : -2.770629e+00 Hartree .. GENERATED FROM PYTHON SOURCE LINES 201-214 Compiling the call ------------------ Everything that changes the shape of the work --- the damping constants, ``cutoff``, the mesh and the spline order --- has to be static. The cell is passed as an argument below rather than captured, so the same compiled function serves variable-cell dynamics as well as fixed-cell: positions, numbers, the cell and the neighbour arrays are all traced. Anything left in the closure --- here the decomposed parameters --- is frozen at trace time, so rebuild the wrapper if you change it. Retracing is driven by shapes, not values, so a trajectory reuses one compiled function as long as the neighbour list keeps its length; its contents may change freely. .. GENERATED FROM PYTHON SOURCE LINES 214-239 .. code-block:: Python jitted = jax.jit( lambda pos, num, box, nl, ptr, sh: fourier_dftd3( pos, num, **damping, fourier_d3_params=params, cell=box, cutoff=cutoff, mesh_dimensions=(32, 32, 32), neighbor_list=nl, neighbor_ptr=ptr, unit_shifts=sh, ) ) compiled_energy, compiled_forces = jitted( positions, numbers, cell, neighbor_list, neighbor_ptr, unit_shifts ) jax.block_until_ready(compiled_energy) print(f"\neager energy : {float(energy[0]):.12f} Hartree") print(f"compiled energy : {float(compiled_energy[0]):.12f} Hartree") print(f"force agreement : {float(jnp.abs(compiled_forces - forces).max()):.2e}") .. rst-class:: sphx-glr-script-out .. code-block:: none eager energy : -0.979196962866 Hartree compiled energy : -0.979196962866 Hartree force agreement : 0.00e+00 .. GENERATED FROM PYTHON SOURCE LINES 240-254 Summary ------- - ``FourierD3Parameters.from_tables`` decomposes the published reference tensor once per species set, on the host, independently of the functional. - ``fourier_dftd3`` is periodic only, needs a coordination-number list whose cutoff equals ``cutoff``, and takes exactly one of ``mesh_dimensions`` or ``mesh_spacing``. - Under ``jax.jit`` the shape-determining arguments must be static; arrays, including the cell, stay traced. One compiled function then serves a trajectory --- fixed-cell or variable-cell --- as long as the input shapes hold. - Forces and the virial are returned explicitly, not obtained by differentiation. For open boundary conditions, or for small systems where the truncation error does not matter, :func:`nvalchemiops.jax.interactions.dispersion.dftd3` remains the right choice. .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 18.623 seconds) .. _sphx_glr_download_examples_dispersion_05_jax_fourier_dftd3_crystal.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 05_jax_fourier_dftd3_crystal.ipynb <05_jax_fourier_dftd3_crystal.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 05_jax_fourier_dftd3_crystal.py <05_jax_fourier_dftd3_crystal.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 05_jax_fourier_dftd3_crystal.zip <05_jax_fourier_dftd3_crystal.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_