Examples#

This page provides practical examples of using torchff to compute common force field energy terms. All examples assume a CUDA-enabled GPU is available.

Each energy module in torchff has two execution paths:

use_customized_ops=True – uses the custom CUDA kernels (fast)
use_customized_ops=False – uses the pure PyTorch reference implementation

Both paths produce the same numerical results and support autograd.

Harmonic Bonds#

Compute the total harmonic bond energy \(E = \sum_i \frac{1}{2} k_i (r_i - r_{0,i})^2\) for a set of bonded atom pairs.

import torch
from torchff.bond import HarmonicBond

device = "cuda"
dtype = torch.float64

# 100 atoms with random 3D coordinates
coords = torch.rand(100, 3, device=device, dtype=dtype, requires_grad=True)

# Define 50 bonds as (atom_i, atom_j) index pairs
bonds = torch.randint(0, 100, (50, 2), device=device)

# Equilibrium bond lengths and force constants
b0 = torch.ones(50, device=device, dtype=dtype) * 0.15   # 0.15 nm
k = torch.ones(50, device=device, dtype=dtype) * 500.0    # kJ/mol/nm^2

# Compute energy using CUDA-accelerated kernels
bond_fn = HarmonicBond(use_customized_ops=True)
energy = bond_fn(coords, bonds, b0, k)
print(f"Bond energy: {energy.item():.4f}")

# Compute forces via autograd
energy.backward()
forces = -coords.grad
print(f"Force on atom 0: {forces[0]}")

Harmonic Angles#

Compute harmonic angle energy \(E = \sum_i \frac{1}{2} k_i (\theta_i - \theta_{0,i})^2\) for triples of atoms (i, j, k) where the angle is measured at the central atom j.

import torch
from torchff.angle import HarmonicAngle

device = "cuda"
dtype = torch.float64

coords = torch.rand(100, 3, device=device, dtype=dtype, requires_grad=True)

# Define 40 angles as (atom_i, atom_j_center, atom_k) triples
angles = torch.randint(0, 100, (40, 3), device=device)

# Equilibrium angles (radians) and force constants
theta0 = torch.ones(40, device=device, dtype=dtype) * 1.911  # ~109.5 degrees
k = torch.ones(40, device=device, dtype=dtype) * 100.0

angle_fn = HarmonicAngle(use_customized_ops=True)
energy = angle_fn(coords, angles, theta0, k)
print(f"Angle energy: {energy.item():.4f}")

energy.backward()
forces = -coords.grad

Periodic Torsions#

Compute periodic torsion (dihedral) energy \(E = \sum_i k_i (1 + \cos(n_i \phi_i - \delta_i))\) for quadruples of atoms defining dihedral angles.

import math
import torch
from torchff.torsion import PeriodicTorsion

device = "cuda"
dtype = torch.float64

coords = torch.rand(100, 3, device=device, dtype=dtype, requires_grad=True)

# Define 30 torsions as (i, j, k, l) quadruples
n_torsions = 30
torsions = torch.randint(0, 100, (n_torsions, 4), device=device)

# Force constants, periodicities, and phase offsets
fc = torch.ones(n_torsions, device=device, dtype=dtype) * 10.0
periodicity = torch.randint(1, 4, (n_torsions,), device=device)
phase = torch.zeros(n_torsions, device=device, dtype=dtype)

torsion_fn = PeriodicTorsion(use_customized_ops=True)
energy = torsion_fn(coords, torsions, fc, periodicity, phase)
print(f"Torsion energy: {energy.item():.4f}")

energy.backward()
forces = -coords.grad

Neighbor List#

Build a neighbor list of atom pairs within a distance cutoff under periodic boundary conditions. This is typically the first step in computing nonbonded interactions.

import torch
import numpy as np
from torchff.nblist import NeighborList

device = "cuda"
dtype = torch.float64

# Simulate a box of particles
n_atoms = 500
box_length = 3.0  # nm
coords = torch.rand(n_atoms, 3, device=device, dtype=dtype) * box_length
box = torch.eye(3, device=device, dtype=dtype) * box_length

cutoff = 1.0  # nm

# Build the neighbor list
nblist = NeighborList(n_atoms, use_customized_ops=True).to(device)
pairs = nblist(coords, box, cutoff)
print(f"Found {pairs.shape[0]} pairs within {cutoff} nm cutoff")
print(f"Pair tensor shape: {pairs.shape}")  # (num_pairs, 2)

Ewald Summation#

Compute long-range electrostatic energy using Ewald summation. torchff supports multipolar Ewald with charges (rank 0), dipoles (rank 1), and quadrupoles (rank 2).

import math
import torch
import numpy as np
from torchff.ewald import Ewald

device = "cuda"
dtype = torch.float64

# Create a system of charged particles in a periodic box
n_atoms = 200
box_length = (n_atoms * 10.0) ** (1.0 / 3.0)

coords = torch.rand(n_atoms, 3, device=device, dtype=dtype, requires_grad=True) * box_length
box = torch.eye(3, device=device, dtype=dtype) * box_length

# Charges (shifted to be charge-neutral)
q_np = np.random.randn(n_atoms)
q_np -= q_np.mean()
q = torch.tensor(q_np, device=device, dtype=dtype, requires_grad=True)

# Ewald parameters
alpha = math.sqrt(-math.log10(2 * 1e-6)) / 9.0
kmax = 10

# Charge-only Ewald (rank=0)
ewald = Ewald(alpha, kmax, rank=0, use_customized_ops=True, return_fields=False)
ewald = ewald.to(device=device, dtype=dtype)

energy = ewald(coords, box, q, p=None, t=None)
print(f"Ewald energy: {energy.item():.6f}")

energy.backward()
forces = -coords.grad

Particle Mesh Ewald (PME)#

PME provides a more efficient \(O(N \log N)\) alternative to direct Ewald summation for large systems. It also supports multipoles up to quadrupole rank.

import torch
import numpy as np
from torchff.pme import PME

device = "cuda"
dtype = torch.float64

n_atoms = 300
box_length = (n_atoms * 10.0) ** (1.0 / 3.0)

coords = torch.rand(n_atoms, 3, device=device, dtype=dtype, requires_grad=True) * box_length
box = torch.eye(3, device=device, dtype=dtype) * box_length

# Charge-neutral charges
q_np = np.random.randn(n_atoms)
q_np -= q_np.mean()
q = torch.tensor(q_np, device=device, dtype=dtype, requires_grad=True)

# Dipole moments
p = torch.randn(n_atoms, 3, device=device, dtype=dtype, requires_grad=True)

alpha = 0.3
max_hkl = 20

# PME with charges and dipoles (rank=1)
pme = PME(alpha, max_hkl, rank=1, use_customized_ops=True)
pme = pme.to(device=device, dtype=dtype)

# PME returns a tuple: (potential, field, field_grad, energy)
result = pme(coords, box, q, p=p, t=None)
energy = result[3]
print(f"PME energy: {energy.item():.6f}")

energy.backward()
forces = -coords.grad

Examples

Contents