nvalchemiops.interactions: Lennard-Jones#

Warp launchers for the Lennard-Jones pair potential. Each launcher accepts a neighbor matrix or a CSR neighbor list. A nonzero switch_width applies the C2 switching function below over the last switch_width of the cutoff.

nvalchemiops.interactions.lj_energy(positions, cell, epsilon, sigma, cutoff, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, fill_value=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, switch_width=0.0, half_neighbor_list=True, device=None)[source]#

Compute Lennard-Jones energies.

Parameters:
  • positions (wp.array, dtype=wp.vec3f or wp.vec3d) – Atomic coordinates. Shape (N,).

  • cell (wp.array, dtype=wp.mat33f or wp.mat33d) – Unit cell matrix. Shape (1,).

  • epsilon (float) – LJ energy parameter (well depth).

  • sigma (float) – LJ length parameter (zero-crossing distance).

  • cutoff (float) – Cutoff distance for interactions.

  • neighbor_matrix (wp.array, shape (N, max_neighbors), dtype=wp.int32, optional) – Neighbor indices in matrix format. Provide either this or neighbor_list.

  • neighbor_matrix_shifts (wp.array, shape (N, max_neighbors), dtype=wp.vec3i, optional) – Periodic shifts for each entry in neighbor_matrix.

  • num_neighbors (wp.array, shape (N,), dtype=wp.int32, optional) – Valid neighbor count per atom; required when using matrix format.

  • fill_value (int, optional) – Sentinel value used to pad neighbor_matrix rows.

  • neighbor_list (wp.array, shape (2, M) or (M,), dtype=wp.int32, optional) – Neighbor target indices in COO/CSR adjacency form; alternative to matrix format.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32, optional) – CSR row pointers; required when neighbor_list is provided.

  • neighbor_shifts (wp.array, shape (M,), dtype=wp.vec3i, optional) – Periodic shifts for each edge in neighbor list format.

  • switch_width (float, optional) – Width of the C2 switching region applied before cutoff. A value of 0.0 (default) disables switching.

  • half_neighbor_list (bool, optional) – True (default) if the neighbor structure contains each pair once. Set to False for full neighbor lists where each pair appears twice.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Per-atom LJ energies (matches input positions dtype). Sum to get total energy.

Return type:

wp.array, shape (N,), dtype=wp.float32 or wp.float64

nvalchemiops.interactions.lj_forces(positions, cell, epsilon, sigma, cutoff, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, fill_value=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, switch_width=0.0, half_neighbor_list=True, device=None)[source]#

Compute Lennard-Jones forces.

Convenience wrapper around nvalchemiops.interactions.lj.lj_energy_forces() that discards the energies and returns only forces.

Parameters:
  • positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic coordinates.

  • cell (wp.array, shape (1,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix.

  • epsilon (float) – LJ energy parameter (well depth).

  • sigma (float) – LJ length parameter (zero-crossing distance).

  • cutoff (float) – Cutoff distance for interactions.

  • neighbor_matrix (wp.array, shape (N, max_neighbors), dtype=wp.int32, optional) – Neighbor indices in matrix format.

  • neighbor_matrix_shifts (wp.array, shape (N, max_neighbors), dtype=wp.vec3i, optional) – Periodic shifts for neighbor matrix format.

  • num_neighbors (wp.array, shape (N,), dtype=wp.int32, optional) – Valid neighbor count per atom (for matrix format).

  • fill_value (int, optional) – Sentinel value used to pad neighbor_matrix.

  • neighbor_list (wp.array, shape (2, M) or (M,), dtype=wp.int32, optional) – Neighbor pairs or target indices in COO/CSR adjacency form.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32, optional) – CSR row pointers; required when neighbor_list is provided.

  • neighbor_shifts (wp.array, shape (M,), dtype=wp.vec3i, optional) – Periodic shifts for neighbor list format.

  • switch_width (float, optional) – Width of the C2 switching region. A value of 0.0 (default) disables switching.

  • half_neighbor_list (bool, optional) – True (default) if the neighbor structure contains each pair once.

  • device (str, optional) – Warp device. If None, inferred from positions.

Returns:

Forces on each atom (matches positions dtype).

Return type:

wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d

nvalchemiops.interactions.lj_energy_forces(positions, cell, epsilon, sigma, cutoff, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, fill_value=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, batch_idx=None, switch_width=0.0, half_neighbor_list=True, device=None, energies_out=None, forces_out=None)[source]#

Compute Lennard-Jones energies and forces.

Parameters:
  • positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic coordinates.

  • cell (wp.array, shape (1,) or (B,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix. Use shape (B,) for batched mode.

  • epsilon (float) – LJ energy parameter (well depth).

  • sigma (float) – LJ length parameter (zero-crossing distance).

  • cutoff (float) – Cutoff distance for interactions.

  • neighbor_matrix (wp.array, shape (N, max_neighbors), dtype=wp.int32, optional) – Neighbor indices in matrix format. Provide either this or neighbor_list.

  • neighbor_matrix_shifts (wp.array, shape (N, max_neighbors), dtype=wp.vec3i, optional) – Periodic shifts for each entry in neighbor_matrix.

  • num_neighbors (wp.array, shape (N,), dtype=wp.int32, optional) – Valid neighbor count per atom; required when using matrix format.

  • fill_value (int, optional) – Sentinel value used to pad neighbor_matrix rows.

  • neighbor_list (wp.array, shape (2, M) or (M,), dtype=wp.int32, optional) – Neighbor target indices in COO/CSR adjacency form; alternative to matrix format.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32, optional) – CSR row pointers; required when neighbor_list is provided.

  • neighbor_shifts (wp.array, shape (M,), dtype=wp.vec3i, optional) – Periodic shifts for each edge in neighbor list format.

  • batch_idx (wp.array, shape (N,), dtype=wp.int32, optional) – System index per atom (0..B-1). Pass None for single-system mode.

  • switch_width (float, optional) – Width of the C2 switching region applied before cutoff. A value of 0.0 (default) disables switching.

  • half_neighbor_list (bool, optional) – True (default) if the neighbor structure contains each pair once. Set to False for full neighbor lists where each pair appears twice.

  • device (str, optional) – Warp device. If None, inferred from positions.

  • energies_out (wp.array, shape (N,), dtype=wp.float32 or wp.float64, optional) – Pre-allocated output buffer for per-atom energies. Modified in-place (zeroed before use). If None, a new array is allocated.

  • forces_out (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d, optional) – Pre-allocated output buffer for forces. Modified in-place (zeroed before use). If None, a new array is allocated.

Returns:

  • wp.array, shape (N,), dtype=wp.float32 or wp.float64 – Per-atom LJ energies (matches input positions dtype).

  • wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d – Forces on each atom (matches positions dtype).

Return type:

tuple[array, array]

nvalchemiops.interactions.lj_energy_forces_virial(positions, cell, epsilon, sigma, cutoff, neighbor_matrix=None, neighbor_matrix_shifts=None, num_neighbors=None, fill_value=None, neighbor_list=None, neighbor_ptr=None, neighbor_shifts=None, batch_idx=None, switch_width=0.0, half_neighbor_list=True, device=None, energies_out=None, forces_out=None, virial_out=None)[source]#

Compute Lennard-Jones energies, forces, and virial tensor.

The virial tensor is needed for pressure/stress calculations in NPT/NPH.

Parameters:
  • positions (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d) – Atomic coordinates.

  • cell (wp.array, shape (1,) or (B,), dtype=wp.mat33f or wp.mat33d) – Unit cell matrix. Use shape (B,) for batched mode.

  • epsilon (float) – LJ energy parameter (well depth).

  • sigma (float) – LJ length parameter (zero-crossing distance).

  • cutoff (float) – Cutoff distance for interactions.

  • neighbor_matrix (wp.array, shape (N, max_neighbors), dtype=wp.int32, optional) – Neighbor indices in matrix format. Provide either this or neighbor_list.

  • neighbor_matrix_shifts (wp.array, shape (N, max_neighbors), dtype=wp.vec3i, optional) – Periodic shifts for each entry in neighbor_matrix.

  • num_neighbors (wp.array, shape (N,), dtype=wp.int32, optional) – Valid neighbor count per atom; required when using matrix format.

  • fill_value (int, optional) – Sentinel value used to pad neighbor_matrix rows.

  • neighbor_list (wp.array, shape (2, M) or (M,), dtype=wp.int32, optional) – Neighbor target indices in COO/CSR adjacency form; alternative to matrix format.

  • neighbor_ptr (wp.array, shape (N+1,), dtype=wp.int32, optional) – CSR row pointers; required when neighbor_list is provided.

  • neighbor_shifts (wp.array, shape (M,), dtype=wp.vec3i, optional) – Periodic shifts for each edge in neighbor list format.

  • batch_idx (wp.array, shape (N,), dtype=wp.int32, optional) – System index per atom (0..B-1). Pass None for single-system mode.

  • switch_width (float, optional) – Width of the C2 switching region applied before cutoff. A value of 0.0 (default) disables switching.

  • half_neighbor_list (bool, optional) – True (default) if the neighbor structure contains each pair once. Set to False for full neighbor lists where each pair appears twice.

  • device (str, optional) – Warp device. If None, inferred from positions.

  • energies_out (wp.array, shape (N,), dtype=wp.float32 or wp.float64, optional) – Pre-allocated output buffer for per-atom energies. Modified in-place (zeroed before use). If None, a new array is allocated.

  • forces_out (wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d, optional) – Pre-allocated output buffer for forces. Modified in-place (zeroed before use). If None, a new array is allocated.

  • virial_out (wp.array, shape (9,) or (B, 9), dtype=wp.float32 or wp.float64, optional) – Pre-allocated output buffer for the virial tensor. Modified in-place (zeroed before use). If None, a new array is allocated.

Returns:

  • wp.array, shape (N,), dtype=wp.float32 or wp.float64 – Per-atom LJ energies (matches input positions dtype).

  • wp.array, shape (N,), dtype=wp.vec3f or wp.vec3d – Forces on each atom (matches positions dtype).

  • wp.array, shape (9,) or (B, 9), dtype=wp.float32 or wp.float64 – Global virial tensor flattened as [xx, xy, xz, yx, yy, yz, zx, zy, zz] (matches input dtype). Shape is (B, 9) in batched mode.

Return type:

tuple[array, array, array]

Switching Function#

nvalchemiops.interactions.switch_c2(r, r_on, r_cut)[source]#

C2 switching function and derivative.

Parameters:
  • r (wp.float64) – Distance.

  • r_on (wp.float64) – Switch-on radius. For r <= r_on, s=1.

  • r_cut (wp.float64) – Cutoff radius. For r >= r_cut, s=0.

Returns:

(s, ds_dr)

Return type:

wp.vec2d