.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/06-splined-potential.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_06-splined-potential.py: .. _example-splined-potential: Splined potentials ================== .. currentmodule:: torchpme :Authors: Michele Ceriotti `@ceriottm `_ This notebook demonstrates the use of the :class:`SplinePotential` class to evaluate potentials for which there is no simple analytical expression for the Fourier-domain filter. .. GENERATED FROM PYTHON SOURCE LINES 17-30 .. code-block:: Python import ase import numpy as np import torch from matplotlib import pyplot as plt import torchpme from torchpme.potentials import CoulombPotential, SplinePotential device = "cpu" dtype = torch.float64 rng = torch.Generator() rng.manual_seed(42) .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 31-36 Defining the potential on a radial grid --------------------------------------- The ``SplinePotential`` can be initialized using simply a pair of arrays: radial positions, and the corresponding values of the potential .. GENERATED FROM PYTHON SOURCE LINES 36-48 .. code-block:: Python x_grid = torch.linspace(0, 10, 8, device=device, dtype=dtype) y_grid = torch.exp(-(x_grid**2) / 2) / torch.pow(torch.pi * 2, torch.tensor([3 / 2])) x_grid_fine = torch.linspace(0, 10, 32, device=device, dtype=dtype) y_grid_fine = torch.exp(-(x_grid_fine**2) / 2) / torch.pow( torch.pi * 2, torch.tensor([3 / 2]) ) spline = SplinePotential(r_grid=x_grid, y_grid=y_grid) spline_fine = SplinePotential(r_grid=x_grid_fine, y_grid=y_grid_fine) .. GENERATED FROM PYTHON SOURCE LINES 49-53 The real-space function can be easily evaluated using the :func:`lr_from_dist ` member function. The convergence with number of spline points is fast, for such a slowly-varying function. .. GENERATED FROM PYTHON SOURCE LINES 53-78 .. code-block:: Python x_test = torch.linspace(0, 10, 256, device=device, dtype=dtype) y_test = torch.exp(-(x_test**2) / 2) / torch.pow(torch.pi * 2, torch.tensor([3 / 2])) y_spline = spline.lr_from_dist(x_test) y_spline_fine = spline_fine.lr_from_dist(x_test) fig, ax = plt.subplots(2, 1, figsize=(4, 3), sharex=True, constrained_layout=True) ax[0].plot(x_grid, y_grid, "b*") ax[0].plot(x_test, y_spline, "b-", label=r"$n_\mathrm{grid}$=8") ax[0].plot(x_grid_fine, y_grid_fine, "r.") ax[0].plot(x_test, y_spline_fine, "r-", label=r"$n_\mathrm{grid}$=32") ax[0].plot(x_test, y_test, "k:") ax[1].set_xlabel(r"$r$ / Å") ax[0].set_ylabel(r"$V$ / a.u.") ax[0].set_xlim(-0.01, 4) ax[0].set_ylim(-0.01, 0.1) ax[1].set_ylabel(r"$\Delta V$ / a.u.") ax[1].plot(x_test, y_spline - y_test, "b-") ax[1].plot(x_test, y_spline_fine - y_test, "r-") ax[1].plot(x_test, y_test * 0, "k:") ax[1].set_ylim(-1e-2, 1e-2) ax[0].legend() .. image-sg:: /examples/images/sphx_glr_06-splined-potential_001.png :alt: 06 splined potential :srcset: /examples/images/sphx_glr_06-splined-potential_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 79-92 Fourier-domain kernel ~~~~~~~~~~~~~~~~~~~~~ A core feature of :class:`SplinePotential ` is that it can evaluate automatically the Fourier-domain kernel that is used in k-space methods. This is done by evaluating .. math :: \hat{f}(k) =4\pi\int \mathrm{d}r \frac{\sin k r}{k} r f(r) in a semin-analytical way - that is, by computing the integral over each segment in the cubic spline. .. GENERATED FROM PYTHON SOURCE LINES 93-111 .. code-block:: Python k_test = torch.linspace(0, 10, 256, device=device, dtype=dtype) yhat_test = torch.exp(-(k_test**2) / 2) # /torch.pow(2*torch.pi,torch.tensor([3/2])) yhat_spline = spline.kernel_from_k_sq(k_test**2) yhat_spline_fine = spline_fine.kernel_from_k_sq(k_test**2) fig, ax = plt.subplots(1, 1, figsize=(4, 3), sharex=True, constrained_layout=True) ax.plot(x_test, yhat_spline, "b-", label=r"$n_\mathrm{grid}$=8") ax.plot(x_test, yhat_spline_fine, "r-", label=r"$n_\mathrm{grid}$=32") ax.plot(k_test, yhat_test, "k:") ax.set_xlabel(r"$k$ / Å$^{-1}$") ax.set_ylabel(r"$\hat{V}$ / a.u.") ax.set_xlim(-0.01, 4) ax.set_ylim(-0.1, 1.1) ax.legend() .. image-sg:: /examples/images/sphx_glr_06-splined-potential_002.png :alt: 06 splined potential :srcset: /examples/images/sphx_glr_06-splined-potential_002.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 112-118 An important consideration is that the splining of the k-space kernel requires a suitable grid. This is usually inferred from the real-space splining, but it is also possible to provided it as a further parameter to the constructor. If the analytical expression for :math:`\hat{V}(k)` is known (and the spline is used for efficiency) one can also provide the values with the parameter ``yhat_grid`` .. GENERATED FROM PYTHON SOURCE LINES 118-135 .. code-block:: Python spline_kgrid = SplinePotential( r_grid=x_grid, y_grid=y_grid, k_grid=torch.linspace(0, 10, 32) ) yhat_spline_kgrid = spline_kgrid.kernel_from_k_sq(k_test**2) fig, ax = plt.subplots(1, 1, figsize=(4, 3), sharex=True, constrained_layout=True) ax.plot(x_test, yhat_spline, "b-", label=r"automatic grid") ax.plot(x_test, yhat_spline_kgrid, "b--", label=r"manual k-space grid") ax.plot(k_test, yhat_test, "k:") ax.set_xlabel(r"$k$ / Å$^{-1}$") ax.set_ylabel(r"$\hat{V}$ / a.u.") ax.set_xlim(-0.01, 4) ax.set_ylim(-0.1, 1.1) ax.legend() .. image-sg:: /examples/images/sphx_glr_06-splined-potential_003.png :alt: 06 splined potential :srcset: /examples/images/sphx_glr_06-splined-potential_003.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 136-145 Reciprocal grid and long-range potentials ----------------------------------------- ``torch-pme`` is all about long-range potentials, and the problem with them is that they converge to zero *very* slowly. In order to address this, :class:`SplinePotential` implements a "reciprocal spline", i.e. the splining grid provided in the definition of the potential is actually defined relative to :math:`1/r`, and "continued" to :math:`1/r\rightarrow 0`. We use a smeared-Coulomb potential to have a more interesting use case. .. GENERATED FROM PYTHON SOURCE LINES 145-154 .. code-block:: Python coulomb = CoulombPotential(smearing=1.0) x_grid = torch.logspace(-2, 2, 100, device=device, dtype=dtype) y_grid = coulomb.lr_from_dist(x_grid) # create a spline potential spline_direct = SplinePotential(r_grid=x_grid, y_grid=y_grid, reciprocal=False) spline = SplinePotential(r_grid=x_grid, y_grid=y_grid, reciprocal=True) .. GENERATED FROM PYTHON SOURCE LINES 155-157 The "direct" spline fails miserably outside the fitting range, while the "reciprocal" spline extrapolates nicely. .. GENERATED FROM PYTHON SOURCE LINES 158-176 .. code-block:: Python t_grid = torch.logspace(-4, 4, 100) z_coul = coulomb.lr_from_dist(t_grid) z_direct = spline_direct.lr_from_dist(t_grid) z_spline = spline.lr_from_dist(t_grid) fig, ax = plt.subplots( 1, 1, figsize=(4, 3), sharey=True, sharex=True, constrained_layout=True ) ax.loglog(t_grid, z_direct, "b-", label="direct spline") ax.loglog(t_grid, z_spline, "r-", label="reciprocal spline") ax.loglog(t_grid, z_coul, "k:") ax.set_xlabel(r"$r$ / Å") ax.set_ylabel(r"$V$ / a.u.") ax.axvspan(1e-2, 1e2, color="gray", alpha=0.3, label="fitted region") ax.set_xlim(1e-4, 1e4) .. image-sg:: /examples/images/sphx_glr_06-splined-potential_004.png :alt: 06 splined potential :srcset: /examples/images/sphx_glr_06-splined-potential_004.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none (0.0001, 10000.0) .. GENERATED FROM PYTHON SOURCE LINES 177-182 This is good, but not magical, and if the tail behavior is not :math:`1/x` the spline will only be able to approximate for a short segment (due to the continuity condition) but will eventually revert to an asymptotic :math:`1/r` behavior. Still, this is much better than a direct spline. .. GENERATED FROM PYTHON SOURCE LINES 182-202 .. code-block:: Python y_grid_2 = y_grid**2 spline_2 = SplinePotential(r_grid=x_grid, y_grid=y_grid_2, reciprocal=True) spline_2_direct = SplinePotential(r_grid=x_grid, y_grid=y_grid_2, reciprocal=False) z_coul_2 = z_coul**2 z_spline_2 = spline_2.lr_from_dist(t_grid) z_spline_2_direct = spline_2_direct.lr_from_dist(t_grid) fig, ax = plt.subplots( 1, 1, figsize=(4, 3), sharey=True, sharex=True, constrained_layout=True ) ax.loglog(t_grid, z_spline_2_direct, "b-", label="direct spline") ax.loglog(t_grid, z_spline_2, "r-", label="reciprocal spline") ax.loglog(t_grid, z_coul_2, "k:") ax.set_xlabel(r"$r$ / Å") ax.set_ylabel(r"$V$ / a.u.") ax.axvspan(1e-2, 1e2, color="gray", alpha=0.3, label="fitted region") ax.set_xlim(1e-4, 1e4) .. image-sg:: /examples/images/sphx_glr_06-splined-potential_005.png :alt: 06 splined potential :srcset: /examples/images/sphx_glr_06-splined-potential_005.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none (0.0001, 10000.0) .. GENERATED FROM PYTHON SOURCE LINES 203-216 Fourier-domain kernel ~~~~~~~~~~~~~~~~~~~~~ The calculation of a Fourier-domain kernel is a very delicate affair for a long-tail potential, which is apparent in the noisy behavior at high-:math:`k`, and the cutoff at the lowest point sampled at small wavevector. These numerical issues can be mitigated, but ultimately have low impact on the accuracy of models built on the splined potential. Note that the initialization of the spline parameters and the calculation of the radial Fourier transform is run in double precision regardless of the type of the input grids. After initialization, further calculations are performed at the level corresponding to the grid precision. .. GENERATED FROM PYTHON SOURCE LINES 217-244 .. code-block:: Python spline = SplinePotential(r_grid=x_grid, y_grid=y_grid, reciprocal=True) x_grid_hiq = torch.logspace(-4, 4, 1000, device=device, dtype=dtype) y_grid_hiq = coulomb.lr_from_dist(x_grid_hiq) spline_hiq = SplinePotential(r_grid=x_grid_hiq, y_grid=y_grid_hiq, reciprocal=True) k_grid = torch.logspace(-4.1, 4, 1000) krn_coul = coulomb.kernel_from_k_sq(k_grid**2) krn_spline = spline.kernel_from_k_sq(k_grid**2) krn_spline_hiq = spline_hiq.kernel_from_k_sq(k_grid**2) fig, ax = plt.subplots( 1, 1, figsize=(4, 3), sharey=True, sharex=True, constrained_layout=True ) ax.loglog(k_grid, krn_spline, "b-", label="low-accuracy spline") ax.loglog(k_grid, krn_spline_hiq, "r-", label="high-accuracy spline") ax.loglog(k_grid, krn_coul, "k:") ax.set_xlabel(r"$k$ / Å$^{-1}$") ax.set_ylabel(r"$\hat{V}$ / a.u.") ax.set_xlim(1e-2, 10) ax.set_ylim(1e-8, 1e4) ax.legend() .. image-sg:: /examples/images/sphx_glr_06-splined-potential_006.png :alt: 06 splined potential :srcset: /examples/images/sphx_glr_06-splined-potential_006.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 245-256 Combining with Fourier filters and meshes ----------------------------------------- To see a potential application of this splining framework, consider the calculation of a "excusion-radius" Coulomb potential, i.e. a smooth Coulomb potential with the short-range region removed (so that) short-range structure does not contribute to the field around an atom. Normally this is achieved through a short-range correction, that however can be cumbersome in some applications (e.g. if one wants to compute explicitly the potential on a grid). .. GENERATED FROM PYTHON SOURCE LINES 256-269 .. code-block:: Python smearing = 0.5 coulomb = CoulombPotential(smearing=smearing, exclusion_radius=None) coulomb_exclude = CoulombPotential(smearing=smearing, exclusion_radius=8.0) x_grid = torch.logspace(-3, 3, 1000) y_grid = coulomb_exclude.lr_from_dist(x_grid) + coulomb_exclude.sr_from_dist(x_grid) # create a spline potential for with the exclusion range built in spline = SplinePotential( r_grid=x_grid, y_grid=y_grid, smearing=smearing, reciprocal=True, yhat_at_zero=0.0 ) .. GENERATED FROM PYTHON SOURCE LINES 270-275 The real-space part of the potential matches the reference and shows clearly how this eliminates the contribution from the atoms within a short-range cutoff. Here we use a very smooth cutoff, but one could as well use a much more aggressive cutoff function. .. GENERATED FROM PYTHON SOURCE LINES 275-294 .. code-block:: Python t_grid = torch.logspace(-3, 3, 1000) y_bare = coulomb.lr_from_dist(t_grid) y_exclude = coulomb_exclude.lr_from_dist(t_grid) + coulomb_exclude.sr_from_dist(t_grid) y_spline = spline.lr_from_dist(t_grid) fig, ax = plt.subplots( 1, 1, figsize=(4, 3), sharey=True, sharex=True, constrained_layout=True ) ax.plot(t_grid, y_spline, "b-", label="8Å exclusion (spline)") ax.plot(t_grid, y_exclude, "r--", label="8Å exclusion (SR)") ax.plot(t_grid, y_bare, "k:", label="smooth Coulomb") ax.set_xlabel(r"$r$ / Å") ax.set_ylabel(r"$V$ / a.u.") ax.set_xlim(0, 20) ax.set_ylim(0, 1.75) ax.legend() .. image-sg:: /examples/images/sphx_glr_06-splined-potential_007.png :alt: 06 splined potential :srcset: /examples/images/sphx_glr_06-splined-potential_007.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 295-296 The k-space kernel has a non-trivial shape .. GENERATED FROM PYTHON SOURCE LINES 296-313 .. code-block:: Python k_grid = torch.logspace(-3, 3, 400) krn_coul = coulomb.kernel_from_k_sq(k_grid**2) krn_spline = spline.kernel_from_k_sq(k_grid**2) fig, ax = plt.subplots( 1, 1, figsize=(4, 3), sharey=True, sharex=True, constrained_layout=True ) ax.semilogx(k_grid, krn_coul, "k:", label="smooth coulomb") ax.semilogx(k_grid, krn_spline, "b-", label=r"8Å exclusion (spline)") ax.set_xlabel(r"$k$ / Å$^{-1}$") ax.set_ylabel(r"$V$ / a.u.") ax.set_xlim(2e-1, 5) ax.set_ylim(-20, 2e2) ax.legend() .. image-sg:: /examples/images/sphx_glr_06-splined-potential_008.png :alt: 06 splined potential :srcset: /examples/images/sphx_glr_06-splined-potential_008.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none .. GENERATED FROM PYTHON SOURCE LINES 314-319 Compute the real-space potential ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ Generate a trial structure -- a distorted rocksalt structure with perturbed positions and charges .. GENERATED FROM PYTHON SOURCE LINES 319-348 .. code-block:: Python structure = ase.Atoms( positions=[ [0, 0, 0], [3, 0, 0], [0, 3, 0], [3, 3, 0], [0, 0, 3], [3, 0, 3], [0, 3, 3], [3, 3, 3], ], cell=[6, 6, 6], symbols="NaClClNaClNaNaCl", ) structure = structure.repeat([3, 3, 1]) displacement = torch.normal(mean=0.0, std=5e-1, size=(len(structure), 3), generator=rng) structure.positions += displacement.numpy() charges = torch.tensor( [[1.0], [-1.0], [-1.0], [1.0], [-1.0], [1.0], [1.0], [-1.0]] * 9, dtype=dtype, device=device, ).reshape(-1, 1) charges += torch.normal(mean=0.0, std=2e-1, size=(len(charges), 1), generator=rng) positions = torch.from_numpy(structure.positions).to(device=device, dtype=dtype) cell = torch.from_numpy(structure.cell.array).to(device=device, dtype=dtype) .. GENERATED FROM PYTHON SOURCE LINES 349-354 We use :class:`MeshInterpolator ` and :class:`KSpaceFilter ` to compute the potential on a grid. Note that the Coulomb potential includes only the k-space part, and therefore has no exclusion zone in reality. .. GENERATED FROM PYTHON SOURCE LINES 354-385 .. code-block:: Python ns = torchpme.lib.kvectors.get_ns_mesh(cell, smearing * 0.5) mesh_interpolator = torchpme.lib.MeshInterpolator( cell=cell, ns_mesh=ns, interpolation_nodes=3, method="P3M" ) kernel_exclusion = torchpme.lib.KSpaceFilter( cell=cell, ns_mesh=ns, kernel=coulomb_exclude, fft_norm="backward", ifft_norm="forward", ) kernel_spline = torchpme.lib.KSpaceFilter( cell=cell, ns_mesh=ns, kernel=spline, fft_norm="backward", ifft_norm="forward", ) mesh_interpolator.compute_weights(positions) rho_mesh = mesh_interpolator.points_to_mesh(particle_weights=charges) ivolume = torch.abs(cell.det()).pow(-1) kernel_exclusion.update(cell, ns) coulomb_mesh = kernel_exclusion.forward(rho_mesh) * ivolume kernel_spline.update(cell, ns) spline_mesh = kernel_spline.forward(rho_mesh) * ivolume .. GENERATED FROM PYTHON SOURCE LINES 386-392 The potential computed using :class:`SplinePotential ` also takes into account the removal of the short-range part of the smooth Coulomb potential, and therefore describes only the slowly-varying part that is generated by the position and charge disorder. .. GENERATED FROM PYTHON SOURCE LINES 392-442 .. code-block:: Python fig, ax = plt.subplots( 1, 2, figsize=(6, 3), sharey=True, sharex=True, constrained_layout=True ) mesh_extent = [ 0, cell[0, 0], 0, cell[1, 1], ] z_plot = coulomb_mesh[0, :, :, 0].cpu().detach().numpy() z_plot = np.vstack([z_plot, z_plot[0, :]]) # Add first row at the bottom z_plot = np.hstack( [z_plot, z_plot[:, 0].reshape(-1, 1)] ) # Add first column at the right z_min, z_max = (z_plot.min(), z_plot.max()) cf_coulomb = ax[0].imshow( z_plot, extent=mesh_extent, vmin=z_min, vmax=z_max, origin="lower", interpolation="bilinear", ) z_plot = spline_mesh[0, :, :, 0].cpu().detach().numpy() z_plot = np.vstack([z_plot, z_plot[0, :]]) # Add first row at the bottom z_plot = np.hstack( [z_plot, z_plot[:, 0].reshape(-1, 1)] ) # Add first column at the right cf_spline = ax[1].imshow( z_plot * 10, extent=mesh_extent, vmin=z_min, vmax=z_max, origin="lower", interpolation="bilinear", ) ax[0].set_title("smooth Coulomb") ax[1].set_title(r"8Å exclusion (spline), 10$\times$") ax[0].set_xlabel(r"$x$ / Å") ax[1].set_xlabel(r"$x$ / Å") ax[0].set_ylabel(r"$y$ / Å") fig.colorbar(cf_coulomb, ax=ax[1], label=r"$V$ / a.u.") fig.show() .. image-sg:: /examples/images/sphx_glr_06-splined-potential_009.png :alt: smooth Coulomb, 8Å exclusion (spline), 10$\times$ :srcset: /examples/images/sphx_glr_06-splined-potential_009.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 2.627 seconds) .. _sphx_glr_download_examples_06-splined-potential.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 06-splined-potential.ipynb <06-splined-potential.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 06-splined-potential.py <06-splined-potential.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 06-splined-potential.zip <06-splined-potential.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_