nonbonded

Convert SMIRNOFF non-bonded parameters into tensors.

Modules:

• smee –

  Differentiably evaluate energies of molecules using SMIRNOFF force fields

Functions:

• convert_nonbonded_handlers –

  Convert a list of SMIRNOFF non-bonded handlers into a tensor potential and

convert_nonbonded_handlers

convert_nonbonded_handlers(
    handlers: list[SMIRNOFFCollection],
    handler_type: str,
    topologies: list[Topology],
    v_site_maps: list[VSiteMap | None],
    parameter_cols: tuple[str, ...],
    attribute_cols: tuple[str, ...] | None = None,
    has_exclusions: bool = True,
) -> tuple[TensorPotential, list[NonbondedParameterMap]]

Convert a list of SMIRNOFF non-bonded handlers into a tensor potential and associated parameter maps.

Notes

This function assumes that all parameters come from the same force field

Parameters:

• handlers (list[SMIRNOFFCollection]) –

  The list of SMIRNOFF non-bonded handlers to convert.

• handler_type (str) –

  The type of non-bonded handler being converted.

• topologies (list[Topology]) –

  The topologies associated with each handler.

• v_site_maps (list[VSiteMap | None]) –

  The virtual site maps associated with each handler.

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

  The ordering of the parameter array columns.

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

  The handler attributes to include in the potential in addition to the intra-molecular scaling factors.

• has_exclusions (bool, default: True ) –

  Whether the handlers are excepted to define exclusions.

Returns:

• tuple[TensorPotential, list[NonbondedParameterMap]] –

  The potential containing tensors of the parameter values, and a list of parameter maps which map the parameters to the interactions they apply to.

Source code in smee/converters/openff/nonbonded.py

def convert_nonbonded_handlers(
    handlers: list[openff.interchange.smirnoff.SMIRNOFFCollection],
    handler_type: str,
    topologies: list[openff.toolkit.Topology],
    v_site_maps: list[smee.VSiteMap | None],
    parameter_cols: tuple[str, ...],
    attribute_cols: tuple[str, ...] | None = None,
    has_exclusions: bool = True,
) -> tuple[smee.TensorPotential, list[smee.NonbondedParameterMap]]:
    """Convert a list of SMIRNOFF non-bonded handlers into a tensor potential and
    associated parameter maps.

    Notes:
        This function assumes that all parameters come from the same force field

    Args:
        handlers: The list of SMIRNOFF non-bonded handlers to convert.
        handler_type: The type of non-bonded handler being converted.
        topologies: The topologies associated with each handler.
        v_site_maps: The virtual site maps associated with each handler.
        parameter_cols: The ordering of the parameter array columns.
        attribute_cols: The handler attributes to include in the potential *in addition*
            to the intra-molecular scaling factors.
        has_exclusions: Whether the handlers are excepted to define exclusions.

    Returns:
        The potential containing tensors of the parameter values, and a list of
        parameter maps which map the parameters to the interactions they apply to.
    """
    attribute_cols = attribute_cols if attribute_cols is not None else []

    assert len(topologies) == len(handlers), "topologies and handlers must match"
    assert len(v_site_maps) == len(handlers), "v-site maps and handlers must match"

    if has_exclusions:
        attribute_cols = (
            "scale_12",
            "scale_13",
            "scale_14",
            "scale_15",
            *attribute_cols,
        )

    potential = smee.converters.openff._openff._handlers_to_potential(
        handlers,
        handler_type,
        parameter_cols,
        attribute_cols,
    )

    parameter_key_to_idx = {
        parameter_key: i for i, parameter_key in enumerate(potential.parameter_keys)
    }
    attribute_to_idx = {column: i for i, column in enumerate(potential.attribute_cols)}

    parameter_maps = []

    for handler, topology, v_site_map in zip(
        handlers, topologies, v_site_maps, strict=True
    ):
        assignment_map = collections.defaultdict(lambda: collections.defaultdict(float))

        n_particles = topology.n_atoms + (
            0 if v_site_map is None else len(v_site_map.keys)
        )

        for topology_key, parameter_key in handler.key_map.items():
            if isinstance(topology_key, openff.interchange.models.VirtualSiteKey):
                continue

            atom_idx = topology_key.atom_indices[0]
            assignment_map[atom_idx][parameter_key_to_idx[parameter_key]] += 1.0

        for topology_key, parameter_key in handler.key_map.items():
            if not isinstance(topology_key, openff.interchange.models.VirtualSiteKey):
                continue

            v_site_idx = v_site_map.key_to_idx[topology_key]

            if parameter_key.associated_handler != "Electrostatics":
                assignment_map[v_site_idx][parameter_key_to_idx[parameter_key]] += 1.0
            else:
                for i, atom_idx in enumerate(topology_key.orientation_atom_indices):
                    mult_key = copy.deepcopy(parameter_key)
                    mult_key.mult = i

                    assignment_map[atom_idx][parameter_key_to_idx[mult_key]] += 1.0
                    assignment_map[v_site_idx][parameter_key_to_idx[mult_key]] += -1.0

        assignment_matrix = torch.zeros(
            (n_particles, len(potential.parameters)), dtype=torch.float64
        )

        for particle_idx in assignment_map:
            for parameter_idx, count in assignment_map[particle_idx].items():
                assignment_matrix[particle_idx, parameter_idx] = count

        if has_exclusions:
            exclusion_to_scale = smee.utils.find_exclusions(topology, v_site_map)
            exclusions = torch.tensor([*exclusion_to_scale])
            exclusion_scale_idxs = torch.tensor(
                [[attribute_to_idx[scale]] for scale in exclusion_to_scale.values()],
                dtype=torch.int64,
            )
        else:
            exclusions = torch.zeros((0, 2), dtype=torch.int64)
            exclusion_scale_idxs = torch.zeros((0, 1), dtype=torch.int64)

        parameter_map = smee.NonbondedParameterMap(
            assignment_matrix=assignment_matrix.to_sparse(),
            exclusions=exclusions,
            exclusion_scale_idxs=exclusion_scale_idxs,
        )
        parameter_maps.append(parameter_map)

    return potential, parameter_maps