Skip to content

nonbonded #

Non-bonded potential energy functions.

Classes:

  • PairwiseDistances –

    A container for the pairwise distances between all particles, possibly within

Functions:

  • compute_pairwise_scales –

    Returns the scale factor for each pair of particles in the system by

  • compute_pairwise –

    Computes all pairwise distances between particles in the system.

  • prepare_lrc_types –

    Finds the unique vdW interactions present in a system, ready to use

  • lorentz_berthelot –

    Computes the Lorentz-Berthelot combination rules for the given parameters.

  • compute_lj_energy –

    Computes the potential energy [kcal / mol] of the vdW interactions using the

  • compute_dexp_energy –

    Compute the potential energy [kcal / mol] of the vdW interactions using the

  • compute_coulomb_energy –

    Computes the potential energy [kcal / mol] of the electrostatic interactions

PairwiseDistances #

Bases: NamedTuple

A container for the pairwise distances between all particles, possibly within a given cutoff.

Attributes:

  • idxs (Tensor) –

    The particle indices of each pair with shape=(n_pairs, 2).

  • deltas (Tensor) –

    The vector between each pair with shape=(n_pairs, 3).

  • distances (Tensor) –

    The distance between each pair with shape=(n_pairs,).

  • cutoff (Tensor | None) –

    The cutoff used when computing the distances.

idxs instance-attribute #

idxs: Tensor

The particle indices of each pair with shape=(n_pairs, 2).

deltas instance-attribute #

deltas: Tensor

The vector between each pair with shape=(n_pairs, 3).

distances instance-attribute #

distances: Tensor

The distance between each pair with shape=(n_pairs,).

cutoff class-attribute instance-attribute #

cutoff: Tensor | None = None

The cutoff used when computing the distances.

compute_pairwise_scales #

compute_pairwise_scales(
    system: TensorSystem, potential: TensorPotential
) -> Tensor

Returns the scale factor for each pair of particles in the system by broadcasting and stacking the exclusions of each topology.

Parameters:

Returns:

  • Tensor –

    The scales for each pair of particles as a flattened upper triangular matrix with shape=(n_particles * (n_particles - 1) / 2,).

Source code in smee/potentials/nonbonded.py 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
def compute_pairwise_scales(
    system: smee.TensorSystem, potential: smee.TensorPotential
) -> torch.Tensor:
    """Returns the scale factor for each pair of particles in the system by
    broadcasting and stacking the exclusions of each topology.

    Args:
        system: The system.
        potential: The potential containing the scale factors to broadcast.

    Returns:
        The scales for each pair of particles as a flattened upper triangular matrix
        with ``shape=(n_particles * (n_particles - 1) / 2,)``.
    """

    n_particles = system.n_particles
    n_pairs = (n_particles * (n_particles - 1)) // 2

    exclusion_idxs, exclusion_scales = _broadcast_exclusions(system, potential)

    pair_scales = smee.utils.ones_like(n_pairs, other=potential.parameters)

    if len(exclusion_idxs) > 0:
        exclusion_idxs, _ = exclusion_idxs.sort(dim=1)  # ensure upper triangle

        pair_idxs = smee.utils.to_upper_tri_idx(
            exclusion_idxs[:, 0], exclusion_idxs[:, 1], n_particles
        )
        pair_scales[pair_idxs] = exclusion_scales

    return pair_scales

compute_pairwise #

compute_pairwise(
    system: TensorSystem,
    conformer: Tensor,
    box_vectors: Tensor | None,
    cutoff: Tensor,
) -> PairwiseDistances

Computes all pairwise distances between particles in the system.

Notes

If the system is not periodic, no cutoff and no PBC will be applied.

Parameters:

  • system (TensorSystem) –

    The system to compute the distances for.

  • conformer (Tensor) –

    The conformer [Å] to evaluate the potential at with shape=(n_confs, n_particles, 3) or shape=(n_particles, 3).

  • box_vectors (Tensor | None) –

    The box vectors [Å] of the system with shape=(n_confs3, 3) or shape=(3, 3) if the system is periodic, or None otherwise.

  • cutoff (Tensor) –

    The cutoff [Å] to apply for periodic systems.

Returns:

  • PairwiseDistances –

    The pairwise distances between each pair of particles within the cutoff.

Source code in smee/potentials/nonbonded.py 
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
def compute_pairwise(
    system: smee.TensorSystem,
    conformer: torch.Tensor,
    box_vectors: torch.Tensor | None,
    cutoff: torch.Tensor,
) -> PairwiseDistances:
    """Computes all pairwise distances between particles in the system.

    Notes:
        If the system is not periodic, no cutoff and no PBC will be applied.

    Args:
        system: The system to compute the distances for.
        conformer: The conformer [Å] to evaluate the potential at with
            ``shape=(n_confs, n_particles, 3)`` or ``shape=(n_particles, 3)``.
        box_vectors: The box vectors [Å] of the system with ``shape=(n_confs3, 3)``
            or ``shape=(3, 3)`` if the system is periodic, or ``None`` otherwise.
        cutoff: The cutoff [Å] to apply for periodic systems.

    Returns:
        The pairwise distances between each pair of particles within the cutoff.
    """
    if system.is_periodic:
        return _compute_pairwise_periodic(conformer, box_vectors, cutoff)
    else:
        return _compute_pairwise_non_periodic(conformer)

prepare_lrc_types #

prepare_lrc_types(
    system: TensorSystem, potential: TensorPotential
) -> tuple[Tensor, Tensor, Tensor]

Finds the unique vdW interactions present in a system, ready to use for computing the long range dispersion correction.

Parameters:

  • system (TensorSystem) –

    The system to prepare the types for.

  • potential (TensorPotential) –

    The potential to prepare the types for.

Returns:

  • tuple[Tensor, Tensor, Tensor] –

    Two tensors containing the i and j indices into potential.paramaters of each unique interaction parameter excluding i==j, the number of ii interactions with shape=(n_params,), and the numbers of ij interactions with shape=(len(idxs_i),).

Source code in smee/potentials/nonbonded.py 
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
def prepare_lrc_types(
    system: smee.TensorSystem,
    potential: smee.TensorPotential,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Finds the unique vdW interactions present in a system, ready to use
    for computing the long range dispersion correction.

    Args:
        system: The system to prepare the types for.
        potential: The potential to prepare the types for.

    Returns:
        Two tensors containing the i and j indices into ``potential.paramaters`` of
        each unique interaction parameter excluding ``i==j``, the number of ``ii``
        interactions with ``shape=(n_params,)``, and the numbers of ``ij`` interactions
        with ``shape=(len(idxs_i),)``.
    """
    n_by_type = collections.defaultdict(int)

    for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
        parameter_counts = topology.parameters["vdW"].assignment_matrix.abs().sum(dim=0)

        for key, count in zip(potential.parameter_keys, parameter_counts, strict=True):
            n_by_type[key] += count.item() * n_copies

    counts = smee.utils.tensor_like(
        [n_by_type[key] for key in potential.parameter_keys], potential.parameters
    )

    n_ii_interactions = (counts * (counts + 1.0)) / 2.0

    idxs_i, idxs_j = torch.triu_indices(len(counts), len(counts), 1)
    n_ij_interactions = counts[idxs_i] * counts[idxs_j]

    idxs_ii = torch.arange(len(counts))

    idxs_i = torch.cat([idxs_i, idxs_ii])
    idxs_j = torch.cat([idxs_j, idxs_ii])

    n_ij_interactions = torch.cat([n_ij_interactions, n_ii_interactions])

    return idxs_i, idxs_j, n_ij_interactions

lorentz_berthelot #

lorentz_berthelot(
    epsilon_a: Tensor,
    epsilon_b: Tensor,
    sigma_a: Tensor,
    sigma_b: Tensor,
) -> tuple[Tensor, Tensor]

Computes the Lorentz-Berthelot combination rules for the given parameters.

Notes

A 'safe' geometric mean is used to avoid NaNs when the parameters are zero. This will yield non-analytic gradients in some cases.

Parameters:

  • epsilon_a (Tensor) –

    The epsilon [kcal / mol] values of the first particle in each pair with shape=(n_pairs, 1).

  • epsilon_b (Tensor) –

    The epsilon [kcal / mol] values of the second particle in each pair with shape=(n_pairs, 1).

  • sigma_a (Tensor) –

    The sigma [kcal / mol] values of the first particle in each pair with shape=(n_pairs, 1).

  • sigma_b (Tensor) –

    The sigma [kcal / mol] values of the second particle in each pair with shape=(n_pairs, 1).

Returns:

  • tuple[Tensor, Tensor] –

    The epsilon [kcal / mol] and sigma [Å] values of each pair, each with shape=(n_pairs, 1).

Source code in smee/potentials/nonbonded.py 
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
def lorentz_berthelot(
    epsilon_a: torch.Tensor,
    epsilon_b: torch.Tensor,
    sigma_a: torch.Tensor,
    sigma_b: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Computes the Lorentz-Berthelot combination rules for the given parameters.

    Notes:
        A 'safe' geometric mean is used to avoid NaNs when the parameters are zero.
        This will yield non-analytic gradients in some cases.

    Args:
        epsilon_a: The epsilon [kcal / mol] values of the first particle in each pair
            with ``shape=(n_pairs, 1)``.
        epsilon_b: The epsilon [kcal / mol] values of the second particle in each pair
            with ``shape=(n_pairs, 1)``.
        sigma_a: The sigma [kcal / mol] values of the first particle in each pair
            with ``shape=(n_pairs, 1)``.
        sigma_b: The sigma [kcal / mol] values of the second particle in each pair
            with ``shape=(n_pairs, 1)``.

    Returns:
        The epsilon [kcal / mol] and sigma [Å] values of each pair, each with
        ``shape=(n_pairs, 1)``.
    """
    return smee.utils.geometric_mean(epsilon_a, epsilon_b), 0.5 * (sigma_a + sigma_b)

compute_lj_energy #

compute_lj_energy(
    system: TensorSystem,
    potential: TensorPotential,
    conformer: Tensor,
    box_vectors: Tensor | None = None,
    pairwise: PairwiseDistances | None = None,
) -> Tensor

Computes the potential energy [kcal / mol] of the vdW interactions using the standard Lennard-Jones potential.

Notes
  • No cutoff / switching function will be applied if the system is not periodic.
  • A switching function will only be applied if the potential has a switch_width attribute.

Parameters:

  • system (TensorSystem) –

    The system to compute the energy for.

  • potential (TensorPotential) –

    The potential energy function to evaluate.

  • conformer (Tensor) –

    The conformer [Å] to evaluate the potential at with shape=(n_confs, n_particles, 3) or shape=(n_particles, 3).

  • box_vectors (Tensor | None, default: None ) –

    The box vectors [Å] of the system with shape=(n_confs, 3, 3) or shape=(3, 3) if the system is periodic, or None otherwise.

  • pairwise (PairwiseDistances | None, default: None ) –

    The pre-computed pairwise distances between each pair of particles in the system. If none, these will be computed within the function.

Returns:

  • Tensor –

    The computed potential energy [kcal / mol] with shape=(n_confs,) if the input conformer has a batch dimension, or shape=() otherwise.

Source code in smee/potentials/nonbonded.py 
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
@smee.potentials.potential_energy_fn(smee.PotentialType.VDW, smee.EnergyFn.VDW_LJ)
def compute_lj_energy(
    system: smee.TensorSystem,
    potential: smee.TensorPotential,
    conformer: torch.Tensor,
    box_vectors: torch.Tensor | None = None,
    pairwise: PairwiseDistances | None = None,
) -> torch.Tensor:
    """Computes the potential energy [kcal / mol] of the vdW interactions using the
    standard Lennard-Jones potential.

    Notes:
        * No cutoff / switching function will be applied if the system is not
          periodic.
        * A switching function will only be applied if the potential has a
          ``switch_width`` attribute.

    Args:
        system: The system to compute the energy for.
        potential: The potential energy function to evaluate.
        conformer: The conformer [Å] to evaluate the potential at with
            ``shape=(n_confs, n_particles, 3)`` or ``shape=(n_particles, 3)``.
        box_vectors: The box vectors [Å] of the system with ``shape=(n_confs, 3, 3)``
            or ``shape=(3, 3)`` if the system is periodic, or ``None`` otherwise.
        pairwise: The pre-computed pairwise distances between each pair of particles
            in the system. If none, these will be computed within the function.

    Returns:
        The computed potential energy [kcal / mol] with ``shape=(n_confs,)`` if the
        input conformer has a batch dimension, or ``shape=()`` otherwise.
    """

    box_vectors = None if not system.is_periodic else box_vectors

    cutoff = potential.attributes[potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)]

    pairwise = (
        pairwise
        if pairwise is not None
        else compute_pairwise(system, conformer, box_vectors, cutoff)
    )

    if system.is_periodic and not torch.isclose(pairwise.cutoff, cutoff):
        raise ValueError("the pairwise cutoff does not match the potential.")

    parameters = smee.potentials.broadcast_parameters(system, potential)
    pair_scales = compute_pairwise_scales(system, potential)

    pairs_1d = smee.utils.to_upper_tri_idx(
        pairwise.idxs[:, 0], pairwise.idxs[:, 1], len(parameters)
    )
    pair_scales = pair_scales[pairs_1d]

    eps_column = potential.parameter_cols.index("epsilon")
    sig_column = potential.parameter_cols.index("sigma")

    eps, sig = lorentz_berthelot(
        parameters[pairwise.idxs[:, 0], eps_column],
        parameters[pairwise.idxs[:, 1], eps_column],
        parameters[pairwise.idxs[:, 0], sig_column],
        parameters[pairwise.idxs[:, 1], sig_column],
    )

    if potential.exceptions is not None:
        exception_idxs, exceptions = smee.potentials.broadcast_exceptions(
            system, potential, pairwise.idxs[:, 0], pairwise.idxs[:, 1]
        )

        eps, sig = eps.clone(), sig.clone()  # prevent in-place modification

        eps[exception_idxs] = exceptions[:, eps_column]
        sig[exception_idxs] = exceptions[:, sig_column]

    x = (sig / pairwise.distances) ** 6
    energies = pair_scales * 4.0 * eps * (x * (x - 1.0))

    if not system.is_periodic:
        return energies.sum(-1)

    switch_fn, switch_width = _compute_switch_fn(potential, pairwise)
    energies *= switch_fn

    energy = energies.sum(-1)
    energy += _compute_lj_lrc(
        system,
        potential.to(precision="double"),
        switch_width.double(),
        pairwise.cutoff.double(),
        torch.det(box_vectors),
    )

    return energy

compute_dexp_energy #

compute_dexp_energy(
    system: TensorSystem,
    potential: TensorPotential,
    conformer: Tensor,
    box_vectors: Tensor | None = None,
    pairwise: PairwiseDistances | None = None,
) -> Tensor

Compute the potential energy [kcal / mol] of the vdW interactions using the double-exponential potential.

Notes
  • No cutoff function will be applied if the system is not periodic.

Parameters:

  • system (TensorSystem) –

    The system to compute the energy for.

  • potential (TensorPotential) –

    The potential energy function to evaluate.

  • conformer (Tensor) –

    The conformer [Å] to evaluate the potential at with shape=(n_confs, n_particles, 3) or shape=(n_particles, 3).

  • box_vectors (Tensor | None, default: None ) –

    The box vectors [Å] of the system with shape=(n_confs, 3, 3) or shape=(3, 3) if the system is periodic, or None otherwise.

  • pairwise (PairwiseDistances | None, default: None ) –

    Pre-computed distances between each pair of particles in the system.

Returns:

  • Tensor –

    The evaluated potential energy [kcal / mol].

Source code in smee/potentials/nonbonded.py 
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
@smee.potentials.potential_energy_fn(smee.PotentialType.VDW, smee.EnergyFn.VDW_DEXP)
def compute_dexp_energy(
    system: smee.TensorSystem,
    potential: smee.TensorPotential,
    conformer: torch.Tensor,
    box_vectors: torch.Tensor | None = None,
    pairwise: PairwiseDistances | None = None,
) -> torch.Tensor:
    """Compute the potential energy [kcal / mol] of the vdW interactions using the
    double-exponential potential.

    Notes:
        * No cutoff function will be applied if the system is not periodic.

    Args:
        system: The system to compute the energy for.
        potential: The potential energy function to evaluate.
        conformer: The conformer [Å] to evaluate the potential at with
            ``shape=(n_confs, n_particles, 3)`` or ``shape=(n_particles, 3)``.
        box_vectors: The box vectors [Å] of the system with ``shape=(n_confs, 3, 3)``
            or ``shape=(3, 3)`` if the system is periodic, or ``None`` otherwise.
        pairwise: Pre-computed distances between each pair of particles
            in the system.

    Returns:
        The evaluated potential energy [kcal / mol].
    """
    box_vectors = None if not system.is_periodic else box_vectors

    cutoff = potential.attributes[potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)]

    pairwise = (
        pairwise
        if pairwise is not None
        else compute_pairwise(system, conformer, box_vectors, cutoff)
    )

    if system.is_periodic and not torch.isclose(pairwise.cutoff, cutoff):
        raise ValueError("the pairwise cutoff does not match the potential.")

    parameters = smee.potentials.broadcast_parameters(system, potential)
    pair_scales = compute_pairwise_scales(system, potential)

    pairs_1d = smee.utils.to_upper_tri_idx(
        pairwise.idxs[:, 0], pairwise.idxs[:, 1], len(parameters)
    )
    pair_scales = pair_scales[pairs_1d]

    eps_column = potential.parameter_cols.index("epsilon")
    r_min_column = potential.parameter_cols.index("r_min")

    eps, r_min = smee.potentials.nonbonded.lorentz_berthelot(
        parameters[pairwise.idxs[:, 0], eps_column],
        parameters[pairwise.idxs[:, 1], eps_column],
        parameters[pairwise.idxs[:, 0], r_min_column],
        parameters[pairwise.idxs[:, 1], r_min_column],
    )

    if potential.exceptions is not None:
        exception_idxs, exceptions = smee.potentials.broadcast_exceptions(
            system, potential, pairwise.idxs[:, 0], pairwise.idxs[:, 1]
        )

        eps, r_min = eps.clone(), r_min.clone()  # prevent in-place modification

        eps[exception_idxs] = exceptions[:, eps_column]
        r_min[exception_idxs] = exceptions[:, r_min_column]

    alpha = potential.attributes[potential.attribute_cols.index("alpha")]
    beta = potential.attributes[potential.attribute_cols.index("beta")]

    x = pairwise.distances / r_min

    energies_repulsion = beta / (alpha - beta) * torch.exp(alpha * (1.0 - x))
    energies_attraction = alpha / (alpha - beta) * torch.exp(beta * (1.0 - x))

    energies = pair_scales * eps * (energies_repulsion - energies_attraction)

    if not system.is_periodic:
        return energies.sum(-1)

    switch_fn, switch_width = _compute_switch_fn(potential, pairwise)
    energies *= switch_fn

    energy = energies.sum(-1)

    energy += _compute_dexp_lrc(
        system,
        potential.to(precision="double"),
        switch_width.double(),
        pairwise.cutoff.double(),
        torch.det(box_vectors),
    )

    return energy

compute_coulomb_energy #

compute_coulomb_energy(
    system: TensorSystem,
    potential: TensorPotential,
    conformer: Tensor,
    box_vectors: Tensor | None = None,
    pairwise: PairwiseDistances | None = None,
) -> Tensor

Computes the potential energy [kcal / mol] of the electrostatic interactions using the Coulomb potential.

Notes
  • No cutoff will be applied if the system is not periodic.
  • PME will be used to compute the energy if the system is periodic.

Parameters:

  • system (TensorSystem) –

    The system to compute the energy for.

  • potential (TensorPotential) –

    The potential energy function to evaluate.

  • conformer (Tensor) –

    The conformer [Å] to evaluate the potential at with shape=(n_confs, n_particles, 3) or shape=(n_particles, 3).

  • box_vectors (Tensor | None, default: None ) –

    The box vectors [Å] of the system with shape=(n_confs, 3, 3) or shape=(3, 3) if the system is periodic, or None otherwise.

  • pairwise (PairwiseDistances | None, default: None ) –

    The pre-computed pairwise distances between each pair of particles in the system. If none, these will be computed within the function.

Returns:

  • Tensor –

    The computed potential energy [kcal / mol] with shape=(n_confs,) if the input conformer has a batch dimension, or shape=() otherwise.

Source code in smee/potentials/nonbonded.py 
 963
 964
 965
 966
 967
 968
 969
 970
 971
 972
 973
 974
 975
 976
 977
 978
 979
 980
 981
 982
 983
 984
 985
 986
 987
 988
 989
 990
 991
 992
 993
 994
 995
 996
 997
 998
 999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
@smee.potentials.potential_energy_fn(
    smee.PotentialType.ELECTROSTATICS, smee.EnergyFn.COULOMB
)
def compute_coulomb_energy(
    system: smee.TensorSystem,
    potential: smee.TensorPotential,
    conformer: torch.Tensor,
    box_vectors: torch.Tensor | None = None,
    pairwise: PairwiseDistances | None = None,
) -> torch.Tensor:
    """Computes the potential energy [kcal / mol] of the electrostatic interactions
    using the Coulomb potential.

    Notes:
        * No cutoff will be applied if the system is not periodic.
        * PME will be used to compute the energy if the system is periodic.

    Args:
        system: The system to compute the energy for.
        potential: The potential energy function to evaluate.
        conformer: The conformer [Å] to evaluate the potential at with
            ``shape=(n_confs, n_particles, 3)`` or ``shape=(n_particles, 3)``.
        box_vectors: The box vectors [Å] of the system with ``shape=(n_confs, 3, 3)``
            or ``shape=(3, 3)`` if the system is periodic, or ``None`` otherwise.
        pairwise: The pre-computed pairwise distances between each pair of particles
            in the system. If none, these will be computed within the function.

    Returns:
        The computed potential energy [kcal / mol] with ``shape=(n_confs,)`` if the
        input conformer has a batch dimension, or ``shape=()`` otherwise.
    """

    box_vectors = None if not system.is_periodic else box_vectors

    cutoff = potential.attributes[potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)]

    pairwise = (
        pairwise
        if pairwise is not None
        else compute_pairwise(system, conformer, box_vectors, cutoff)
    )

    if system.is_periodic and not torch.isclose(pairwise.cutoff, cutoff):
        raise ValueError("the distance cutoff does not match the potential.")

    if potential.exceptions is not None:
        raise NotImplementedError("exceptions are not supported for charges.")

    if system.is_periodic:
        return _compute_coulomb_energy_periodic(
            system, conformer, box_vectors, potential, pairwise
        )
    else:
        return _compute_coulomb_energy_non_periodic(system, potential, pairwise)