.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/electrostatics/01_coulomb_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_electrostatics_01_coulomb_example.py: Coulomb Electrostatic Interactions ================================== This example demonstrates how to compute Coulomb electrostatic interactions using GPU-accelerated kernels. The implementation supports both direct (undamped) Coulomb interactions and damped interactions used as the real-space component in Ewald/PME methods. In this example you will learn: - How to compute Coulomb energies and forces between charged particles - The difference between undamped (1/r) and damped (erfc(αr)/r) interactions - Using both neighbor list (COO) and neighbor matrix formats - Computing forces via autograd vs explicit force kernels - Batch evaluation for multiple systems - Computing gradients with respect to charges .. 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 35-39 Setup and Imports ----------------- First, we import the necessary modules. The Coulomb API provides three main functions: ``coulomb_energy``, ``coulomb_forces``, and ``coulomb_energy_forces``. .. GENERATED FROM PYTHON SOURCE LINES 39-49 .. code-block:: Python from __future__ import annotations import torch from nvalchemiops.interactions.electrostatics.coulomb import ( coulomb_energy, coulomb_energy_forces, ) .. GENERATED FROM PYTHON SOURCE LINES 50-53 Configure Device ---------------- We use CUDA if available for GPU acceleration. .. GENERATED FROM PYTHON SOURCE LINES 53-61 .. code-block:: Python if torch.cuda.is_available(): device = torch.device("cuda:0") print(f"Using CUDA device: {torch.cuda.get_device_name(0)}") else: device = torch.device("cpu") print("Using CPU") .. rst-class:: sphx-glr-script-out .. code-block:: none Using CUDA device: NVIDIA L4 .. GENERATED FROM PYTHON SOURCE LINES 62-77 Basic Coulomb Calculation ------------------------- We start with a simple two-atom system: a +1 and -1 charge separated by 2.0 Å. This demonstrates both undamped Coulomb (direct 1/r) and damped Coulomb (erfc(αr)/r used in Ewald/PME real-space). The undamped Coulomb energy for a pair is: .. math:: E = \frac{q_i q_j}{r} The damped Coulomb (Ewald real-space) energy is: .. math:: E = \frac{q_i q_j \, \mathrm{erfc}(\alpha r)}{r} .. GENERATED FROM PYTHON SOURCE LINES 77-107 .. code-block:: Python # Two charges: +1 and -1 separated by 2.0 Angstroms positions = torch.tensor( [ [0.0, 0.0, 0.0], # Charge +1 [2.0, 0.0, 0.0], # Charge -1 ], dtype=torch.float64, device=device, ) charges = torch.tensor([1.0, -1.0], dtype=torch.float64, device=device) # Large cell to avoid periodic images cell = torch.tensor( [[[100.0, 0.0, 0.0], [0.0, 100.0, 0.0], [0.0, 0.0, 100.0]]], dtype=torch.float64, device=device, ) # Direct neighbor specification: pair (0, 1) neighbor_list = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32, device=device) neighbor_ptr = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) neighbor_shifts = torch.zeros((2, 3), dtype=torch.int32, device=device) print("System setup:") print(f" Positions: {positions.cpu().numpy()}") print(f" Charges: {charges.cpu().numpy()}") print(" Distance: 2.0 Å") .. rst-class:: sphx-glr-script-out .. code-block:: none System setup: Positions: [[0. 0. 0.] [2. 0. 0.]] Charges: [ 1. -1.] Distance: 2.0 Å .. GENERATED FROM PYTHON SOURCE LINES 108-111 Undamped Coulomb ~~~~~~~~~~~~~~~~ With α=0, we get the direct Coulomb interaction. .. GENERATED FROM PYTHON SOURCE LINES 111-133 .. code-block:: Python energies_undamped, forces_undamped = coulomb_energy_forces( positions, charges, cell, cutoff=10.0, alpha=0.0, # Undamped neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr, neighbor_shifts=neighbor_shifts, ) r = 2.0 expected_energy = 1.0 * (-1.0) / r # = -0.5 print("\nUndamped Coulomb (α=0):") print( f" Total energy: {energies_undamped.sum().item():.6f} (expected: {expected_energy:.6f})" ) print(f" Forces on atom 0: {forces_undamped[0].cpu().numpy()}") print(f" Forces on atom 1: {forces_undamped[1].cpu().numpy()}") .. rst-class:: sphx-glr-script-out .. code-block:: none Undamped Coulomb (α=0): Total energy: -0.500000 (expected: -0.500000) Forces on atom 0: [0.25 0. 0. ] Forces on atom 1: [-0.25 0. 0. ] .. GENERATED FROM PYTHON SOURCE LINES 134-138 Damped Coulomb (Ewald Real-Space) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ With α > 0, the interaction is damped by erfc(αr). This is the short-range component used in Ewald/PME methods. .. GENERATED FROM PYTHON SOURCE LINES 138-155 .. code-block:: Python energies_damped, forces_damped = coulomb_energy_forces( positions, charges, cell, cutoff=10.0, alpha=0.3, # Damped with α=0.3 neighbor_list=neighbor_list, neighbor_ptr=neighbor_ptr, neighbor_shifts=neighbor_shifts, ) print("\nDamped Coulomb (α=0.3):") print(f" Total energy: {energies_damped.sum().item():.6f}") print(f" Forces on atom 0: {forces_damped[0].cpu().numpy()}") print(f" Forces on atom 1: {forces_damped[1].cpu().numpy()}") .. rst-class:: sphx-glr-script-out .. code-block:: none Damped Coulomb (α=0.3): Total energy: -0.198072 Forces on atom 0: [0.21712248 0. 0. ] Forces on atom 1: [-0.21712248 0. 0. ] .. GENERATED FROM PYTHON SOURCE LINES 156-159 Newton's Third Law Check ~~~~~~~~~~~~~~~~~~~~~~~~ The sum of forces should be zero for an isolated system. .. GENERATED FROM PYTHON SOURCE LINES 159-165 .. code-block:: Python net_force = forces_undamped.sum(dim=0) print("\nNewton's 3rd law check:") print(f" Net force: {net_force.cpu().numpy()}") print(f" Magnitude: {net_force.norm().item():.2e} (should be ~0)") .. rst-class:: sphx-glr-script-out .. code-block:: none Newton's 3rd law check: Net force: [0. 0. 0.] Magnitude: 0.00e+00 (should be ~0) .. GENERATED FROM PYTHON SOURCE LINES 166-170 Autograd vs Explicit Forces --------------------------- The Coulomb functions support automatic differentiation. Forces computed via autograd should match the explicit forces returned by the kernel. .. GENERATED FROM PYTHON SOURCE LINES 170-199 .. code-block:: Python positions_grad = torch.tensor( [ [0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [1.0, 1.73, 0.0], # Equilateral triangle ], dtype=torch.float64, device=device, requires_grad=True, ) charges_3 = torch.tensor([1.0, -0.5, -0.5], dtype=torch.float64, device=device) cell_3 = torch.tensor( [[[50.0, 0.0, 0.0], [0.0, 50.0, 0.0], [0.0, 0.0, 50.0]]], dtype=torch.float64, device=device, ) # All pairs as neighbors neighbor_list_3 = torch.tensor( [[0, 0, 1, 1, 2, 2], [1, 2, 0, 2, 0, 1]], dtype=torch.int32, device=device ) neighbor_ptr_3 = torch.tensor([0, 2, 4, 6], dtype=torch.int32, device=device) neighbor_shifts_3 = torch.zeros((6, 3), dtype=torch.int32, device=device) print("\nThree-atom triangular system:") print(f" Charges: {charges_3.cpu().numpy()}") .. rst-class:: sphx-glr-script-out .. code-block:: none Three-atom triangular system: Charges: [ 1. -0.5 -0.5] .. GENERATED FROM PYTHON SOURCE LINES 200-201 Get explicit forces from the kernel: .. GENERATED FROM PYTHON SOURCE LINES 201-217 .. code-block:: Python _, explicit_forces = coulomb_energy_forces( positions_grad.detach(), charges_3, cell_3, cutoff=10.0, alpha=0.0, neighbor_list=neighbor_list_3, neighbor_ptr=neighbor_ptr_3, neighbor_shifts=neighbor_shifts_3, ) print("\nExplicit Forces:") for i, f in enumerate(explicit_forces): print(f" Atom {i}: {f.cpu().numpy()}") .. rst-class:: sphx-glr-script-out .. code-block:: none Explicit Forces: Atom 0: [0.18766678 0.10841352 0. ] Atom 1: [-0.09366661 -0.05420676 0. ] Atom 2: [-0.09400016 -0.05420676 0. ] .. GENERATED FROM PYTHON SOURCE LINES 218-219 Get autograd forces via backpropagation: .. GENERATED FROM PYTHON SOURCE LINES 219-243 .. code-block:: Python energies = coulomb_energy( positions_grad, charges_3, cell_3, cutoff=10.0, alpha=0.0, neighbor_list=neighbor_list_3, neighbor_ptr=neighbor_ptr_3, neighbor_shifts=neighbor_shifts_3, ) total_energy = energies.sum() total_energy.backward() autograd_forces = -positions_grad.grad print(f"\nTotal energy: {total_energy.item():.6f}") print("\nAutograd Forces (-∂E/∂r):") for i, f in enumerate(autograd_forces): print(f" Atom {i}: {f.cpu().numpy()}") force_diff = (autograd_forces - explicit_forces).abs().max().item() print(f"\nMax force difference: {force_diff:.2e}") print(f"Forces match: {'✓' if force_diff < 1e-8 else '✗'}") .. rst-class:: sphx-glr-script-out .. code-block:: none Total energy: -0.375111 Autograd Forces (-∂E/∂r): Atom 0: [ 0.18766678 0.10841352 -0. ] Atom 1: [-0.09366661 -0.05420676 -0. ] Atom 2: [-0.09400016 -0.05420676 -0. ] Max force difference: 0.00e+00 Forces match: ✓ .. GENERATED FROM PYTHON SOURCE LINES 244-252 Neighbor List vs Neighbor Matrix Format --------------------------------------- The Coulomb functions support two neighbor formats: - **Neighbor List (COO)**: Shape ``(2, num_pairs)`` where each pair is listed once - **Neighbor Matrix**: Shape ``(N, max_neighbors)`` with padding Both formats produce identical results. .. GENERATED FROM PYTHON SOURCE LINES 252-274 .. code-block:: Python # Create a 4-atom system in a square positions_4 = torch.tensor( [ [0.0, 0.0, 0.0], [2.0, 0.0, 0.0], [0.0, 2.0, 0.0], [2.0, 2.0, 0.0], ], dtype=torch.float64, device=device, ) charges_4 = torch.tensor([1.0, -1.0, -1.0, 1.0], dtype=torch.float64, device=device) cell_4 = torch.tensor( [[[50.0, 0.0, 0.0], [0.0, 50.0, 0.0], [0.0, 0.0, 50.0]]], dtype=torch.float64, device=device, ) print("\n4-atom square system:") print(f" Charges: {charges_4.cpu().numpy()}") .. rst-class:: sphx-glr-script-out .. code-block:: none 4-atom square system: Charges: [ 1. -1. -1. 1.] .. GENERATED FROM PYTHON SOURCE LINES 275-276 Neighbor list format (COO): .. GENERATED FROM PYTHON SOURCE LINES 276-298 .. code-block:: Python neighbor_list_coo = torch.tensor( [[0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3], [1, 2, 3, 0, 2, 3, 0, 1, 3, 0, 1, 2]], dtype=torch.int32, device=device, ) neighbor_ptr_coo = torch.tensor([0, 3, 6, 9, 12], dtype=torch.int32, device=device) neighbor_shifts_coo = torch.zeros((12, 3), dtype=torch.int32, device=device) energies_list, forces_list = coulomb_energy_forces( positions_4, charges_4, cell_4, cutoff=10.0, alpha=0.0, neighbor_list=neighbor_list_coo, neighbor_ptr=neighbor_ptr_coo, neighbor_shifts=neighbor_shifts_coo, ) print(f"\nNeighbor List format: Total energy = {energies_list.sum().item():.6f}") .. rst-class:: sphx-glr-script-out .. code-block:: none Neighbor List format: Total energy = -1.292893 .. GENERATED FROM PYTHON SOURCE LINES 299-300 Neighbor matrix format: .. GENERATED FROM PYTHON SOURCE LINES 300-333 .. code-block:: Python fill_value = 4 neighbor_matrix = torch.tensor( [ [1, 2, 3, fill_value], [0, 2, 3, fill_value], [0, 1, 3, fill_value], [0, 1, 2, fill_value], ], dtype=torch.int32, device=device, ) neighbor_matrix_shifts = torch.zeros((4, 4, 3), dtype=torch.int32, device=device) energies_matrix, forces_matrix = coulomb_energy_forces( positions_4, charges_4, cell_4, cutoff=10.0, alpha=0.0, neighbor_matrix=neighbor_matrix, neighbor_matrix_shifts=neighbor_matrix_shifts, fill_value=fill_value, ) print(f"Neighbor Matrix format: Total energy = {energies_matrix.sum().item():.6f}") energy_diff = (energies_list.sum() - energies_matrix.sum()).abs().item() force_diff = (forces_list - forces_matrix).abs().max().item() print(f"\nEnergy difference: {energy_diff:.2e}") print(f"Max force difference: {force_diff:.2e}") print(f"Formats match: {'✓' if energy_diff < 1e-10 and force_diff < 1e-10 else '✗'}") .. rst-class:: sphx-glr-script-out .. code-block:: none Neighbor Matrix format: Total energy = -1.292893 Energy difference: 0.00e+00 Max force difference: 0.00e+00 Formats match: ✓ .. GENERATED FROM PYTHON SOURCE LINES 334-338 Batch Evaluation ---------------- Multiple independent systems can be evaluated simultaneously using the ``batch_idx`` parameter. Each atom is assigned to a system index. .. GENERATED FROM PYTHON SOURCE LINES 338-363 .. code-block:: Python # Create two identical 2-atom systems system_positions = torch.tensor( [[0.0, 0.0, 0.0], [2.5, 0.0, 0.0]], dtype=torch.float64, device=device ) system_charges = torch.tensor([1.0, -1.0], dtype=torch.float64, device=device) # Concatenate for batch (second system offset by 100 Å) positions_batch = torch.cat([system_positions, system_positions + 100.0]) charges_batch = torch.cat([system_charges, system_charges]) batch_idx = torch.tensor([0, 0, 1, 1], dtype=torch.int32, device=device) cell_batch = torch.tensor( [[[200.0, 0.0, 0.0], [0.0, 200.0, 0.0], [0.0, 0.0, 200.0]]], dtype=torch.float64, device=device, ) # Neighbor pairs: (0,1) and (2,3) nl_batch = torch.tensor([[0, 1, 2, 3], [1, 0, 3, 2]], dtype=torch.int32, device=device) ptr_batch = torch.tensor([0, 1, 2, 3, 4], dtype=torch.int32, device=device) shifts_batch = torch.zeros((4, 3), dtype=torch.int32, device=device) print("\nBatch Evaluation: Two identical 2-atom systems") .. rst-class:: sphx-glr-script-out .. code-block:: none Batch Evaluation: Two identical 2-atom systems .. GENERATED FROM PYTHON SOURCE LINES 364-365 Batch calculation: .. GENERATED FROM PYTHON SOURCE LINES 365-397 .. code-block:: Python energies_batch, _ = coulomb_energy_forces( positions_batch, charges_batch, cell_batch, cutoff=10.0, alpha=0.0, neighbor_list=nl_batch, neighbor_ptr=ptr_batch, neighbor_shifts=shifts_batch, batch_idx=batch_idx, ) # Single system for comparison single_nl = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32, device=device) single_ptr = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) single_shifts = torch.zeros((1, 3), dtype=torch.int32, device=device) energies_single, _ = coulomb_energy_forces( system_positions, system_charges, cell_batch, cutoff=10.0, alpha=0.0, neighbor_list=single_nl, neighbor_ptr=single_ptr, neighbor_shifts=single_shifts, ) print(f"Single system energy: {energies_single.sum().item():.6f}") print(f"Batch system 0 energy: {energies_batch[:2].sum().item():.6f}") print(f"Batch system 1 energy: {energies_batch[2:].sum().item():.6f}") .. rst-class:: sphx-glr-script-out .. code-block:: none Single system energy: -0.200000 Batch system 0 energy: -0.400000 Batch system 1 energy: -0.400000 .. GENERATED FROM PYTHON SOURCE LINES 398-402 Charge Gradients ---------------- The autograd support extends to charge gradients, enabling sensitivity analysis and optimization with respect to charges. .. GENERATED FROM PYTHON SOURCE LINES 402-435 .. code-block:: Python positions_q = torch.tensor( [[0.0, 0.0, 0.0], [3.0, 0.0, 0.0]], dtype=torch.float64, device=device ) charges_q = torch.tensor( [1.0, -1.0], dtype=torch.float64, device=device, requires_grad=True ) cell_q = torch.tensor( [[[100.0, 0.0, 0.0], [0.0, 100.0, 0.0], [0.0, 0.0, 100.0]]], dtype=torch.float64, device=device, ) nl_q = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32, device=device) ptr_q = torch.tensor([0, 1, 2], dtype=torch.int32, device=device) shifts_q = torch.zeros((1, 3), dtype=torch.int32, device=device) energies_q = coulomb_energy( positions_q, charges_q, cell_q, cutoff=10.0, alpha=0.0, neighbor_list=nl_q, neighbor_ptr=ptr_q, neighbor_shifts=shifts_q, ) total_energy_q = energies_q.sum() total_energy_q.backward() print("\nCharge Gradients") print(f"Two atoms at r=3.0 Å, charges: {charges_q.data.cpu().numpy()}") print(f"Total energy: {total_energy_q.item():.6f}") .. rst-class:: sphx-glr-script-out .. code-block:: none Charge Gradients Two atoms at r=3.0 Å, charges: [ 1. -1.] Total energy: -0.166667 .. GENERATED FROM PYTHON SOURCE LINES 436-442 Verify against analytical result: For E = q₀q₁/r: - ∂E/∂q₀ = q₁/r = -1/3 ≈ -0.333 - ∂E/∂q₁ = q₀/r = 1/3 ≈ 0.333 .. GENERATED FROM PYTHON SOURCE LINES 442-450 .. code-block:: Python print(f"\n∂E/∂q₀ = {charges_q.grad[0].item():.6f} (expected: {-1.0 / 3.0:.6f})") print(f"∂E/∂q₁ = {charges_q.grad[1].item():.6f} (expected: {1.0 / 3.0:.6f})") match_q0 = abs(charges_q.grad[0].item() - (-1.0 / 3.0)) < 1e-6 match_q1 = abs(charges_q.grad[1].item() - (1.0 / 3.0)) < 1e-6 print(f"Match: {'✓' if match_q0 and match_q1 else '✗'}") .. rst-class:: sphx-glr-script-out .. code-block:: none ∂E/∂q₀ = -0.166667 (expected: -0.333333) ∂E/∂q₁ = 0.166667 (expected: 0.333333) Match: ✗ .. GENERATED FROM PYTHON SOURCE LINES 451-464 Summary ------- This example demonstrated: 1. **Undamped Coulomb**: Direct 1/r interactions with ``alpha=0`` 2. **Damped Coulomb**: Ewald real-space with ``alpha > 0`` 3. **Autograd support**: Forces via ``-positions.grad`` 4. **Neighbor formats**: Both COO list and matrix formats 5. **Batch evaluation**: Multiple systems with ``batch_idx`` 6. **Charge gradients**: Sensitivity via ``charges.grad`` For visualization of energy/force profiles as a function of distance, see the separate visualization scripts in the examples folder. .. GENERATED FROM PYTHON SOURCE LINES 464-466 .. code-block:: Python print("\nCoulomb example complete!") .. rst-class:: sphx-glr-script-out .. code-block:: none Coulomb example complete! .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 2.610 seconds) .. _sphx_glr_download_examples_electrostatics_01_coulomb_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_coulomb_example.ipynb <01_coulomb_example.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 01_coulomb_example.py <01_coulomb_example.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 01_coulomb_example.zip <01_coulomb_example.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_