smee

smee

Differentiably evaluate energies of molecules using SMIRNOFF force fields

Modules:

• converters –

  Convert to / from smee tensor representations.

• geometry –

  Compute internal coordinates (e.g. bond lengths).

• mm –

  Compute differentiable ensemble averages using OpenMM and SMEE.

• potentials –

  Evaluate the potential energy of parameterized topologies.

• tests –
• utils –

  General utility functions

Classes:

• EnergyFn –

  An enumeration of the energy functions supported by smee out of the box.

• PotentialType –

  An enumeration of the potential types supported by smee out of the box.

• NonbondedParameterMap –

  A map between atom indices part of a particular valence interaction (e.g.

• TensorConstraints –

  A tensor representation of a set of distance constraints between pairs of

• TensorForceField –

  A tensor representation of a SMIRNOFF force field.

• TensorPotential –

  A tensor representation of a valence SMIRNOFF parameter handler

• TensorSystem –

  A tensor representation of a 'full' system.

• TensorTopology –

  A tensor representation of a molecular topology that has been assigned force

• TensorVSites –

  A tensor representation of a set of virtual sites parameters.

• ValenceParameterMap –

  A map between atom indices part of a particular valence interaction (e.g.

• VSiteMap –

  A map between virtual sites that have been added to a topology and their

Functions:

• add_v_site_coords –

  Appends the coordinates of any virtual sites to a conformer (or batch of

• compute_v_site_coords –

  Computes the positions of a set of virtual sites relative to a specified

• compute_energy –

  Computes the potential energy [kcal / mol] of a system / topology in a given

• compute_energy_potential –

  Computes the potential energy [kcal / mol] due to a SMIRNOFF potential

Attributes:

• CUTOFF_ATTRIBUTE –

  The attribute that should be used to store the cutoff distance of a potential.

• SWITCH_ATTRIBUTE –

  The attribute that should be used to store the switch width of a potential, if the

CUTOFF_ATTRIBUTE module-attribute

CUTOFF_ATTRIBUTE = 'cutoff'

The attribute that should be used to store the cutoff distance of a potential.

SWITCH_ATTRIBUTE module-attribute

SWITCH_ATTRIBUTE = 'switch_width'

The attribute that should be used to store the switch width of a potential, if the potential should use the standard OpenMM switch function.

This attribute should be omitted if the potential should not use a switch function.

EnergyFn

Bases: _StrEnum

An enumeration of the energy functions supported by smee out of the box.

PotentialType

Bases: _StrEnum

An enumeration of the potential types supported by smee out of the box.

NonbondedParameterMap dataclass

NonbondedParameterMap(
    assignment_matrix: Tensor,
    exclusions: Tensor,
    exclusion_scale_idxs: Tensor,
)

A map between atom indices part of a particular valence interaction (e.g. torsion indices) and the corresponding parameter in a TensorPotential

Methods:

• to –

  Cast this object to the specified device.

Attributes:

• assignment_matrix (Tensor) –

  A sparse tensor that yields the parameters assigned to each particle in the

• exclusions (Tensor) –

  Indices of pairs of particles (i.e. atoms or virtual sites) that should

• exclusion_scale_idxs (Tensor) –

  Indices into the tensor of handler attributes defining the 1-n scaling factors

assignment_matrix instance-attribute

assignment_matrix: Tensor

A sparse tensor that yields the parameters assigned to each particle in the system when multiplied with the corresponding handler parameters, with shape=(n_particles, n_parameters).

exclusions instance-attribute

exclusions: Tensor

Indices of pairs of particles (i.e. atoms or virtual sites) that should have their interactions scaled by some factor with shape=(n_exclusions, 2).

exclusion_scale_idxs instance-attribute

exclusion_scale_idxs: Tensor

Indices into the tensor of handler attributes defining the 1-n scaling factors with shape=(n_exclusions, 1).

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> NonbondedParameterMap

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "NonbondedParameterMap":
    """Cast this object to the specified device."""
    return NonbondedParameterMap(
        _cast(self.assignment_matrix, device, precision),
        _cast(self.exclusions, device, precision),
        _cast(self.exclusion_scale_idxs, device, precision),
    )

TensorConstraints dataclass

TensorConstraints(idxs: Tensor, distances: Tensor)

A tensor representation of a set of distance constraints between pairs of atoms.

Methods:

• to –

  Cast this object to the specified device.

Attributes:

• idxs (Tensor) –

  The indices of the atoms involved in each constraint with

• distances (Tensor) –

  The distance [Å] between each pair of atoms with shape=(n_constraints,)

idxs instance-attribute

idxs: Tensor

The indices of the atoms involved in each constraint with shape=(n_constraints, 2)

distances instance-attribute

distances: Tensor

The distance [Å] between each pair of atoms with shape=(n_constraints,)

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> TensorConstraints

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "TensorConstraints":
    """Cast this object to the specified device."""
    return TensorConstraints(
        _cast(self.idxs, device, precision),
        _cast(self.distances, device, precision),
    )

TensorForceField dataclass

TensorForceField(
    potentials: list[TensorPotential],
    v_sites: TensorVSites | None = None,
)

A tensor representation of a SMIRNOFF force field.

Methods:

• to –

  Cast this object to the specified device.

Attributes:

• potentials (list[TensorPotential]) –

  The terms and associated parameters of the potential energy function.

• v_sites (TensorVSites | None) –

  Parameters used to add and define the coords of v-sites in the system. The

potentials instance-attribute

potentials: list[TensorPotential]

The terms and associated parameters of the potential energy function.

v_sites class-attribute instance-attribute

v_sites: TensorVSites | None = None

Parameters used to add and define the coords of v-sites in the system. The non-bonded parameters of any v-sites are stored in relevant potentials, e.g. 'vdW' or 'Electrostatics'.

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> TensorForceField

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "TensorForceField":
    """Cast this object to the specified device."""
    return TensorForceField(
        [potential.to(device, precision) for potential in self.potentials],
        None if self.v_sites is None else self.v_sites.to(device, precision),
    )

TensorPotential dataclass

TensorPotential(
    type: str,
    fn: str,
    parameters: Tensor,
    parameter_keys: list[PotentialKey],
    parameter_cols: tuple[str, ...],
    parameter_units: tuple[Unit, ...],
    attributes: Tensor | None = None,
    attribute_cols: tuple[str, ...] | None = None,
    attribute_units: tuple[Unit, ...] | None = None,
    exceptions: dict[tuple[int, int], int] | None = None,
)

A tensor representation of a valence SMIRNOFF parameter handler

Methods:

• to –

  Cast this object to the specified device.

Attributes:

• type (str) –

  The type of handler associated with these parameters

• fn (str) –

  The associated potential energy function

• parameters (Tensor) –

  The values of the parameters with shape=(n_parameters, n_parameter_cols)

• parameter_keys (list[PotentialKey]) –

  Unique keys associated with each parameter with length=(n_parameters)

• parameter_cols (tuple[str, ...]) –

  The names of each column of parameters.

• parameter_units (tuple[Unit, ...]) –

  The units of each parameter in parameters.

• attributes (Tensor | None) –

  The attributes defined on a handler such as 1-4 scaling factors with

• attribute_cols (tuple[str, ...] | None) –

  The names of each column of attributes.

• attribute_units (tuple[Unit, ...] | None) –

  The units of each attribute in attributes.

• exceptions (dict[tuple[int, int], int] | None) –

  A lookup for custom cross-interaction parameters that should override any mixing

type instance-attribute

type: str

The type of handler associated with these parameters

fn instance-attribute

fn: str

The associated potential energy function

parameters instance-attribute

parameters: Tensor

The values of the parameters with shape=(n_parameters, n_parameter_cols)

parameter_keys instance-attribute

parameter_keys: list[PotentialKey]

Unique keys associated with each parameter with length=(n_parameters)

parameter_cols instance-attribute

parameter_cols: tuple[str, ...]

The names of each column of parameters.

parameter_units instance-attribute

parameter_units: tuple[Unit, ...]

The units of each parameter in parameters.

attributes class-attribute instance-attribute

attributes: Tensor | None = None

The attributes defined on a handler such as 1-4 scaling factors with shape=(n_attribute_cols,)

attribute_cols class-attribute instance-attribute

attribute_cols: tuple[str, ...] | None = None

The names of each column of attributes.

attribute_units class-attribute instance-attribute

attribute_units: tuple[Unit, ...] | None = None

The units of each attribute in attributes.

exceptions class-attribute instance-attribute

exceptions: dict[tuple[int, int], int] | None = None

A lookup for custom cross-interaction parameters that should override any mixing rules.

Each key should correspond to the indices of the two parameters whose mixing rule should be overridden, and each value the index of the parameter that contains the 'pre-mixed' parameter to use instead.

For now, all exceptions are assumed to be symmetric, i.e. if (a, b) is an exception then (b, a) is also an exception, and so only one of the two should be defined.

As a note of caution, not all potentials (e.g. common valence potentials) support such exceptions, and these are predominantly useful for non-bonded potentials.

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> TensorPotential

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "TensorPotential":
    """Cast this object to the specified device."""
    return TensorPotential(
        self.type,
        self.fn,
        _cast(self.parameters, device, precision),
        self.parameter_keys,
        self.parameter_cols,
        self.parameter_units,
        (
            None
            if self.attributes is None
            else _cast(self.attributes, device, precision)
        ),
        self.attribute_cols,
        self.attribute_units,
        self.exceptions,
    )

TensorSystem dataclass

TensorSystem(
    topologies: list[TensorTopology],
    n_copies: list[int],
    is_periodic: bool,
)

A tensor representation of a 'full' system.

Methods:

• to –

  Cast this object to the specified device.

Attributes:

• topologies (list[TensorTopology]) –

  The topologies of the individual molecules in the system.

• n_copies (list[int]) –

  The number of copies of each topology to include in the system.

• is_periodic (bool) –

  Whether the system is periodic or not.

• n_atoms (int) –

  The number of atoms in the system.

• n_v_sites (int) –

  The number of v-sites in the system.

• n_particles (int) –

  The number of atoms + v-sites in the system.

topologies instance-attribute

topologies: list[TensorTopology]

The topologies of the individual molecules in the system.

n_copies instance-attribute

n_copies: list[int]

The number of copies of each topology to include in the system.

is_periodic instance-attribute

is_periodic: bool

Whether the system is periodic or not.

n_atoms property

n_atoms: int

The number of atoms in the system.

n_v_sites property

n_v_sites: int

The number of v-sites in the system.

n_particles property

n_particles: int

The number of atoms + v-sites in the system.

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> TensorSystem

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "TensorSystem":
    """Cast this object to the specified device."""
    return TensorSystem(
        [topology.to(device, precision) for topology in self.topologies],
        self.n_copies,
        self.is_periodic,
    )

TensorTopology dataclass

TensorTopology(
    atomic_nums: Tensor,
    formal_charges: Tensor,
    bond_idxs: Tensor,
    bond_orders: Tensor,
    parameters: dict[str, ParameterMap],
    v_sites: VSiteMap | None = None,
    constraints: TensorConstraints | None = None,
    residue_idxs: list[int] | None = None,
    residue_ids: list[str] | None = None,
    chain_idxs: list[int] | None = None,
    chain_ids: list[str] | None = None,
)

A tensor representation of a molecular topology that has been assigned force field parameters.

Methods:

• to –

  Cast this object to the specified device.

Attributes:

• atomic_nums (Tensor) –

  The atomic numbers of each atom in the topology with shape=(n_atoms,)

• formal_charges (Tensor) –

  The formal charge of each atom in the topology with shape=(n_atoms,)

• bond_idxs (Tensor) –

  The indices of the atoms involved in each bond with shape=(n_bonds, 2)

• bond_orders (Tensor) –

  The bond orders of each bond with shape=(n_bonds,)

• parameters (dict[str, ParameterMap]) –

  The parameters that have been assigned to the topology.

• v_sites (VSiteMap | None) –

  The v-sites that have been assigned to the topology.

• constraints (TensorConstraints | None) –

  Distance constraints that should be applied during MD simulations. These

• residue_idxs (list[int] | None) –

  The index of the residue that each atom in the topology belongs to with

• residue_ids (list[str] | None) –

  The names of the residues that each atom belongs to with length=n_residues.

• chain_idxs (list[int] | None) –

  The index of the chain that each atom in the topology belongs to with

• chain_ids (list[str] | None) –

  The names of the chains that each atom belongs to with length=n_chains.

• n_atoms (int) –

  The number of atoms in the topology.

• n_bonds (int) –

  The number of bonds in the topology.

• n_residues (int) –

  The number of residues in the topology

• n_chains (int) –

  The number of chains in the topology

• n_v_sites (int) –

  The number of v-sites in the topology.

• n_particles (int) –

  The number of atoms + v-sites in the topology.

atomic_nums instance-attribute

atomic_nums: Tensor

The atomic numbers of each atom in the topology with shape=(n_atoms,)

formal_charges instance-attribute

formal_charges: Tensor

The formal charge of each atom in the topology with shape=(n_atoms,)

bond_idxs instance-attribute

bond_idxs: Tensor

The indices of the atoms involved in each bond with shape=(n_bonds, 2)

bond_orders instance-attribute

bond_orders: Tensor

The bond orders of each bond with shape=(n_bonds,)

parameters instance-attribute

parameters: dict[str, ParameterMap]

The parameters that have been assigned to the topology.

v_sites class-attribute instance-attribute

v_sites: VSiteMap | None = None

The v-sites that have been assigned to the topology.

constraints class-attribute instance-attribute

constraints: TensorConstraints | None = None

Distance constraints that should be applied during MD simulations. These will not be used outside of MD simulations.

residue_idxs class-attribute instance-attribute

residue_idxs: list[int] | None = None

The index of the residue that each atom in the topology belongs to with length=n_atoms.

residue_ids class-attribute instance-attribute

residue_ids: list[str] | None = None

The names of the residues that each atom belongs to with length=n_residues.

chain_idxs class-attribute instance-attribute

chain_idxs: list[int] | None = None

The index of the chain that each atom in the topology belongs to with length=n_atoms.

chain_ids class-attribute instance-attribute

chain_ids: list[str] | None = None

The names of the chains that each atom belongs to with length=n_chains.

n_atoms property

n_atoms: int

The number of atoms in the topology.

n_bonds property

n_bonds: int

The number of bonds in the topology.

n_residues property

n_residues: int

The number of residues in the topology

n_chains property

n_chains: int

The number of chains in the topology

n_v_sites property

n_v_sites: int

The number of v-sites in the topology.

n_particles property

n_particles: int

The number of atoms + v-sites in the topology.

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> TensorTopology

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "TensorTopology":
    """Cast this object to the specified device."""
    return TensorTopology(
        self.atomic_nums,
        self.formal_charges,
        self.bond_idxs,
        self.bond_orders,
        {k: v.to(device, precision) for k, v in self.parameters.items()},
        None if self.v_sites is None else self.v_sites.to(device, precision),
        (
            None
            if self.constraints is None
            else self.constraints.to(device, precision)
        ),
        self.residue_idxs,
        self.residue_ids,
        self.chain_ids,
    )

TensorVSites dataclass

TensorVSites(
    keys: List[VirtualSiteKey],
    weights: list[Tensor],
    parameters: Tensor,
)

A tensor representation of a set of virtual sites parameters.

Methods:

• default_units –

  The default units of each v-site parameter.

• to –

  Cast this object to the specified device.

Attributes:

• keys (List[VirtualSiteKey]) –

  The unique keys associated with each v-site with length=(n_v_sites)

• weights (list[Tensor]) –

  A matrix of weights that, when applied to the 'orientiational' atoms, yields a

• parameters (Tensor) –

  The distance, in-plane and out-of-plane angles with shape=(n_v_sites, 3)

• parameter_units (dict[str, Unit]) –

  The units of each v-site parameter.

keys instance-attribute

keys: List[VirtualSiteKey]

The unique keys associated with each v-site with length=(n_v_sites)

weights instance-attribute

weights: list[Tensor]

A matrix of weights that, when applied to the 'orientiational' atoms, yields a basis that the virtual site coordinate parameters can be projected onto with shape=(n_v_sites, 3, 3)

parameters instance-attribute

parameters: Tensor

The distance, in-plane and out-of-plane angles with shape=(n_v_sites, 3)

parameter_units property

parameter_units: dict[str, Unit]

The units of each v-site parameter.

default_units classmethod

default_units() -> dict[str, Unit]

The default units of each v-site parameter.

Source code in smee/_models.py

@classmethod
def default_units(cls) -> dict[str, openff.units.Unit]:
    """The default units of each v-site parameter."""
    return {
        "distance": _ANGSTROM,
        "inPlaneAngle": _RADIANS,
        "outOfPlaneAngle": _RADIANS,
    }

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> TensorVSites

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "TensorVSites":
    """Cast this object to the specified device."""
    return TensorVSites(
        self.keys,
        [_cast(weight, device, precision) for weight in self.weights],
        _cast(self.parameters, device, precision),
    )

ValenceParameterMap dataclass

ValenceParameterMap(
    particle_idxs: Tensor, assignment_matrix: Tensor
)

A map between atom indices part of a particular valence interaction (e.g. torsion indices) and the corresponding parameter in a TensorPotential

Methods:

• to –

  Cast this object to the specified device.

Attributes:

• particle_idxs (Tensor) –

  The indices of the particles (e.g. atoms or virtual sites) involved in an

• assignment_matrix (Tensor) –

  A sparse tensor that yields the assigned parameters when multiplied with the

particle_idxs instance-attribute

particle_idxs: Tensor

The indices of the particles (e.g. atoms or virtual sites) involved in an interaction with shape=(n_interactions, n_cols). For a bond n_cols=2, for angles n_cols=3 etc.

assignment_matrix instance-attribute

assignment_matrix: Tensor

A sparse tensor that yields the assigned parameters when multiplied with the corresponding handler parameters, with shape=(n_interacting, n_parameters).

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> ValenceParameterMap

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "ValenceParameterMap":
    """Cast this object to the specified device."""
    return ValenceParameterMap(
        _cast(self.particle_idxs, device, precision),
        _cast(self.assignment_matrix, device, precision),
    )

VSiteMap dataclass

VSiteMap(
    keys: list[VirtualSiteKey],
    key_to_idx: dict[VirtualSiteKey, int],
    parameter_idxs: Tensor,
)

A map between virtual sites that have been added to a topology and their corresponding 'parameters' used to position them.

Methods:

• to –

  Cast this object to the specified device.

Attributes:

• keys (list[VirtualSiteKey]) –

  The keys used to identify each v-site.

• key_to_idx (dict[VirtualSiteKey, int]) –

  A map between the unique keys associated with each v-site and their index in

• parameter_idxs (Tensor) –

  The indices of the corresponding v-site parameters with shape=(n_v_sites, 1)

keys instance-attribute

keys: list[VirtualSiteKey]

The keys used to identify each v-site.

key_to_idx instance-attribute

key_to_idx: dict[VirtualSiteKey, int]

A map between the unique keys associated with each v-site and their index in the topology

parameter_idxs instance-attribute

parameter_idxs: Tensor

The indices of the corresponding v-site parameters with shape=(n_v_sites, 1)

to

to(
    device: DeviceType | None = None,
    precision: Precision | None = None,
) -> VSiteMap

Cast this object to the specified device.

Source code in smee/_models.py

def to(
    self, device: DeviceType | None = None, precision: Precision | None = None
) -> "VSiteMap":
    """Cast this object to the specified device."""
    return VSiteMap(
        self.keys, self.key_to_idx, _cast(self.parameter_idxs, device, precision)
    )

add_v_site_coords

add_v_site_coords(
    v_sites: VSiteMap,
    conformer: Tensor,
    force_field: TensorForceField,
) -> Tensor

Appends the coordinates of any virtual sites to a conformer (or batch of conformers) containing only atomic coordinates.

Notes

• This function only supports appending v-sites to the end of the list of coordinates, and not interleaving them between existing atomic coordinates.

Parameters:

• v_sites (VSiteMap) –

  A mapping between the virtual sites to add and their corresponding force field parameters.

• conformer (Tensor) –

  The conformer(s) to add the virtual sites to with shape=(n_atoms, 3) or shape=(n_batches, n_atoms, 3) and units of [Å].

• force_field (TensorForceField) –

  The force field containing the virtual site parameters.

Returns:

• Tensor –

  The full conformer(s) with both atomic and virtual site coordinates [Å] with shape=(n_atoms+n_v_sites, 3) or shape=(n_batches, n_atoms+n_v_sites, 3).

Source code in smee/geometry.py

def add_v_site_coords(
    v_sites: "smee.VSiteMap",
    conformer: torch.Tensor,
    force_field: "smee.TensorForceField",
) -> torch.Tensor:
    """Appends the coordinates of any virtual sites to a conformer (or batch of
    conformers) containing only atomic coordinates.

    Notes:
        * This function only supports appending v-sites to the end of the list of
          coordinates, and not interleaving them between existing atomic coordinates.

    Args:
        v_sites: A mapping between the virtual sites to add and their corresponding
            force field parameters.
        conformer: The conformer(s) to add the virtual sites to with
            ``shape=(n_atoms, 3)`` or ``shape=(n_batches, n_atoms, 3)`` and units of
            [Å].
        force_field: The force field containing the virtual site parameters.

    Returns:
        The full conformer(s) with both atomic and virtual site coordinates [Å] with
        ``shape=(n_atoms+n_v_sites, 3)`` or ``shape=(n_batches, n_atoms+n_v_sites, 3)``.
    """

    v_site_coords = compute_v_site_coords(v_sites, conformer, force_field)

    return torch.cat([conformer, v_site_coords], dim=(1 if conformer.ndim == 3 else 0))

compute_v_site_coords

compute_v_site_coords(
    v_sites: VSiteMap,
    conformer: Tensor,
    force_field: TensorForceField,
) -> Tensor

Computes the positions of a set of virtual sites relative to a specified conformer or batch of conformers.

Parameters:

• v_sites (VSiteMap) –

  A mapping between the virtual sites to add and their corresponding force field parameters.

• conformer (Tensor) –

  The conformer(s) to add the virtual sites to with shape=(n_atoms, 3) or shape=(n_batches, n_atoms, 3) and units of [Å].

• force_field (TensorForceField) –

  The force field containing the virtual site parameters.

Returns:

• Tensor –

  A tensor of virtual site positions [Å] with shape=(n_v_sites, 3) or shape=(n_batches, n_v_sites, 3).

Source code in smee/geometry.py

def compute_v_site_coords(
    v_sites: "smee.VSiteMap",
    conformer: torch.Tensor,
    force_field: "smee.TensorForceField",
) -> torch.Tensor:
    """Computes the positions of a set of virtual sites relative to a specified
    conformer or batch of conformers.

    Args:
        v_sites: A mapping between the virtual sites to add and their corresponding
            force field parameters.
        conformer: The conformer(s) to add the virtual sites to with
            ``shape=(n_atoms, 3)`` or ``shape=(n_batches, n_atoms, 3)`` and units of
            [Å].
        force_field: The force field containing the virtual site parameters.

    Returns:
        A tensor of virtual site positions [Å] with ``shape=(n_v_sites, 3)`` or
        ``shape=(n_batches, n_v_sites, 3)``.
    """

    is_batched = conformer.ndim == 3

    if not is_batched:
        conformer = torch.unsqueeze(conformer, 0)

    if len(v_sites.parameter_idxs) > 0:
        local_frame_coords = force_field.v_sites.parameters[v_sites.parameter_idxs]
        local_coord_frames = _build_v_site_coord_frames(v_sites, conformer, force_field)

        v_site_coords = _convert_v_site_coords(local_frame_coords, local_coord_frames)
    else:
        v_site_coords = smee.utils.zeros_like((len(conformer), 0, 3), other=conformer)

    if not is_batched:
        v_site_coords = torch.squeeze(v_site_coords, 0)

    return v_site_coords

compute_energy

compute_energy(
    system: TensorSystem | TensorTopology,
    force_field: TensorForceField,
    conformer: Tensor,
    box_vectors: Tensor | None = None,
) -> Tensor

Computes the potential energy [kcal / mol] of a system / topology in a given conformation(s).

Parameters:

• system (TensorSystem | TensorTopology) –

  The system or topology to compute the potential energy of.

• force_field (TensorForceField) –

  The force field that defines the potential energy function.

• conformer (Tensor) –

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

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

  The box vectors of the system with shape=(3, 3) if the system is periodic, or None if the system is non-periodic.

Returns:

• Tensor –

  The potential energy of the conformer(s) [kcal / mol].

Source code in smee/potentials/_potentials.py

def compute_energy(
    system: smee.TensorSystem | smee.TensorTopology,
    force_field: smee.TensorForceField,
    conformer: torch.Tensor,
    box_vectors: torch.Tensor | None = None,
) -> torch.Tensor:
    """Computes the potential energy [kcal / mol] of a system / topology in a given
    conformation(s).

    Args:
        system: The system or topology to compute the potential energy of.
        force_field: The force field that defines the potential energy function.
        conformer: The conformer(s) to evaluate the potential at with
            ``shape=(n_particles, 3)`` or ``shape=(n_confs, n_particles, 3)``.
        box_vectors: The box vectors of the system with ``shape=(3, 3)`` if the
            system is periodic, or ``None`` if the system is non-periodic.

    Returns:
        The potential energy of the conformer(s) [kcal / mol].
    """

    # register the built-in potential energy functions
    importlib.import_module("smee.potentials.nonbonded")
    importlib.import_module("smee.potentials.valence")

    system, conformer, box_vectors = _prepare_inputs(system, conformer, box_vectors)
    pairwise = _precompute_pairwise(system, force_field, conformer, box_vectors)

    energy = smee.utils.zeros_like(
        conformer.shape[0] if conformer.ndim == 3 else 1, conformer
    )

    for potential in force_field.potentials:
        energy += compute_energy_potential(
            system, potential, conformer, box_vectors, pairwise
        )

    return energy

compute_energy_potential

compute_energy_potential(
    system: TensorSystem | TensorTopology,
    potential: TensorPotential,
    conformer: Tensor,
    box_vectors: Tensor | None = None,
    pairwise: Optional[PairwiseDistances] = None,
) -> Tensor

Computes the potential energy [kcal / mol] due to a SMIRNOFF potential handler for a given conformer(s).

Parameters:

• system (TensorSystem | TensorTopology) –

  The system or topology to compute the potential energy of.

• potential (TensorPotential) –

  The potential describing the energy function to compute.

• conformer (Tensor) –

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

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

  The box vectors of the system with shape=(3, 3) or shape=(n_confs, 3, 3)if the system is periodic, orNone`` if the system is non-periodic.

• pairwise (Optional[PairwiseDistances], default: None ) –

  Pre-computed pairwise distances between particles in the system.

Returns:

• Tensor –

  The potential energy of the conformer(s) [kcal / mol].

Source code in smee/potentials/_potentials.py

def compute_energy_potential(
    system: smee.TensorSystem | smee.TensorTopology,
    potential: smee.TensorPotential,
    conformer: torch.Tensor,
    box_vectors: torch.Tensor | None = None,
    pairwise: typing.Optional["smee.potentials.nonbonded.PairwiseDistances"] = None,
) -> torch.Tensor:
    """Computes the potential energy [kcal / mol] due to a SMIRNOFF potential
    handler for a given conformer(s).

    Args:
        system: The system or topology to compute the potential energy of.
        potential: The potential describing the energy function to compute.
        conformer: The conformer(s) to evaluate the potential at with
            ``shape=(n_particles, 3)`` or ``shape=(n_confs, n_particles, 3)``.
        box_vectors: The box vectors of the system with ``shape=(3, 3)`` or
            shape=(n_confs, 3, 3)`` if the system is periodic, or ``None`` if the system
            is non-periodic.
        pairwise: Pre-computed pairwise distances between particles in the system.

    Returns:
        The potential energy of the conformer(s) [kcal / mol].
    """

    # register the built-in potential energy functions
    importlib.import_module("smee.potentials.nonbonded")
    importlib.import_module("smee.potentials.valence")

    system, conformer, box_vectors = _prepare_inputs(system, conformer, box_vectors)

    energy_fn = _POTENTIAL_ENERGY_FUNCTIONS[(potential.type, potential.fn)]
    energy_fn_spec = inspect.signature(energy_fn)

    energy_fn_kwargs = {}

    if "box_vectors" in energy_fn_spec.parameters:
        energy_fn_kwargs["box_vectors"] = box_vectors
    if "pairwise" in energy_fn_spec.parameters:
        energy_fn_kwargs["pairwise"] = pairwise

    return energy_fn(system, potential, conformer, **energy_fn_kwargs)