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.

Setup and Imports

First, we import the necessary modules. The Coulomb API provides three main functions: coulomb_energy, coulomb_forces, and coulomb_energy_forces.

from __future__ import annotations

import torch

from nvalchemiops.interactions.electrostatics.coulomb import (
    coulomb_energy,
    coulomb_energy_forces,
)

Configure Device

We use CUDA if available for GPU acceleration.

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")

Using CUDA device: NVIDIA L4

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:

\[E = \frac{q_i q_j}{r}\]

The damped Coulomb (Ewald real-space) energy is:

\[E = \frac{q_i q_j \, \mathrm{erfc}(\alpha r)}{r}\]

# 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 Å")

System setup:
  Positions: [[0. 0. 0.]
 [2. 0. 0.]]
  Charges: [ 1. -1.]
  Distance: 2.0 Å

Undamped Coulomb

With α=0, we get the direct Coulomb interaction.

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()}")

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.  ]

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.

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()}")

Damped Coulomb (α=0.3):
  Total energy: -0.198072
  Forces on atom 0: [0.21712248 0.         0.        ]
  Forces on atom 1: [-0.21712248  0.          0.        ]

Newton’s Third Law Check

The sum of forces should be zero for an isolated system.

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)")

Newton's 3rd law check:
  Net force: [0. 0. 0.]
  Magnitude: 0.00e+00 (should be ~0)

Autograd vs Explicit Forces

The Coulomb functions support automatic differentiation. Forces computed via autograd should match the explicit forces returned by the kernel.

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()}")

Three-atom triangular system:
  Charges: [ 1.  -0.5 -0.5]

Get explicit forces from the kernel:

_, 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()}")

Explicit Forces:
  Atom 0: [0.18766678 0.10841352 0.        ]
  Atom 1: [-0.09366661 -0.05420676  0.        ]
  Atom 2: [-0.09400016 -0.05420676  0.        ]

Get autograd forces via backpropagation:

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 '✗'}")

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: ✓

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.

# 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()}")

4-atom square system:
  Charges: [ 1. -1. -1.  1.]

Neighbor list format (COO):

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}")

Neighbor List format: Total energy = -1.292893

Neighbor matrix format:

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 '✗'}")

Neighbor Matrix format: Total energy = -1.292893

Energy difference: 0.00e+00
Max force difference: 0.00e+00
Formats match: ✓

Batch Evaluation

Multiple independent systems can be evaluated simultaneously using the batch_idx parameter. Each atom is assigned to a system index.

# 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")

Batch Evaluation: Two identical 2-atom systems

Batch calculation:

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}")

Single system energy: -0.200000
Batch system 0 energy: -0.400000
Batch system 1 energy: -0.400000

Charge Gradients

The autograd support extends to charge gradients, enabling sensitivity analysis and optimization with respect to charges.

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}")

Charge Gradients
Two atoms at r=3.0 Å, charges: [ 1. -1.]
Total energy: -0.166667

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

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 '✗'}")

∂E/∂q₀ = -0.166667 (expected: -0.333333)
∂E/∂q₁ = 0.166667 (expected: 0.333333)
Match: ✗

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.

print("\nCoulomb example complete!")

Coulomb example complete!