openff

Tensor representations of SMIRNOFF force fields.

Modules:

• nonbonded –

  Convert SMIRNOFF non-bonded parameters into tensors.

• valence –

  Convert SMIRNOFF valence parameters into tensors.

Functions:

• convert_handlers –

  Convert a set of SMIRNOFF parameter handlers into a set of tensor potentials.

• convert_interchange –

  Convert a list of interchange objects into tensor potentials.

• smirnoff_parameter_converter –

  A decorator used to flag a function as being able to convert a parameter handlers

convert_handlers

convert_handlers(
    handlers: list[SMIRNOFFCollection],
    topologies: list[Topology],
    v_site_maps: list[VSiteMap | None] | None = None,
    potentials: (
        list[tuple[TensorPotential, list[ParameterMap]]]
        | None
    ) = None,
    constraints: (
        list[TensorConstraints | None] | None
    ) = None,
) -> list[tuple[TensorPotential, list[ParameterMap]]]

Convert a set of SMIRNOFF parameter handlers into a set of tensor potentials.

Parameters:

• handlers (list[SMIRNOFFCollection]) –

  The SMIRNOFF parameter handler collections for a set of interchange objects to convert.

• topologies (list[Topology]) –

  The topologies associated with each interchange object.

• v_site_maps (list[VSiteMap | None] | None, default: None ) –

  The v-site maps associated with each interchange object.

• constraints (list[TensorConstraints | None] | None, default: None ) –

  Any distance constraints between atoms.

• potentials (list[tuple[TensorPotential, list[ParameterMap]]] | None, default: None ) –

  Already converted parameter handlers that may be required as dependencies.

Returns:

• list[tuple[TensorPotential, list[ParameterMap]]] –

  The potential containing the values of the parameters in each handler collection, and a list of maps (one per topology) between molecule elements (e.g. bond indices) and parameter indices.

Examples:

>>> from openff.toolkit import ForceField, Molecule
>>> from openff.interchange import Interchange
>>>
>>> force_field = ForceField("openff_unconstrained-2.0.0.offxml")
>>> molecules = [Molecule.from_smiles("CCO"), Molecule.from_smiles("CC")]
>>>
>>> interchanges = [
...     Interchange.from_smirnoff(force_field, molecule.to_topology())
...     for molecule in molecules
... ]
>>> vdw_handlers = [
...     interchange.collections["vdW"] for interchange in interchanges
... ]
>>>
>>> vdw_potential, applied_vdw_parameters = convert_handlers(interchanges)

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

def convert_handlers(
    handlers: list[openff.interchange.smirnoff.SMIRNOFFCollection],
    topologies: list[openff.toolkit.Topology],
    v_site_maps: list[smee.VSiteMap | None] | None = None,
    potentials: (
        list[tuple[smee.TensorPotential, list[smee.ParameterMap]]] | None
    ) = None,
    constraints: list[smee.TensorConstraints | None] | None = None,
) -> list[tuple[smee.TensorPotential, list[smee.ParameterMap]]]:
    """Convert a set of SMIRNOFF parameter handlers into a set of tensor potentials.

    Args:
        handlers: The SMIRNOFF parameter handler collections for a set of interchange
            objects to convert.
        topologies: The topologies associated with each interchange object.
        v_site_maps: The v-site maps associated with each interchange object.
        constraints: Any distance constraints between atoms.
        potentials: Already converted parameter handlers that may be required as
            dependencies.

    Returns:
        The potential containing the values of the parameters in each handler
        collection, and a list of maps (one per topology) between molecule elements
        (e.g. bond indices) and parameter indices.

    Examples:

        >>> from openff.toolkit import ForceField, Molecule
        >>> from openff.interchange import Interchange
        >>>
        >>> force_field = ForceField("openff_unconstrained-2.0.0.offxml")
        >>> molecules = [Molecule.from_smiles("CCO"), Molecule.from_smiles("CC")]
        >>>
        >>> interchanges = [
        ...     Interchange.from_smirnoff(force_field, molecule.to_topology())
        ...     for molecule in molecules
        ... ]
        >>> vdw_handlers = [
        ...     interchange.collections["vdW"] for interchange in interchanges
        ... ]
        >>>
        >>> vdw_potential, applied_vdw_parameters = convert_handlers(interchanges)
    """
    importlib.import_module("smee.converters.openff.nonbonded")
    importlib.import_module("smee.converters.openff.valence")

    handler_types = {handler.type for handler in handlers}
    assert len(handler_types) == 1, "multiple handler types found"
    handler_type = next(iter(handler_types))

    assert len(handlers) == len(topologies), "mismatched number of topologies"

    if handler_type not in _CONVERTERS:
        raise NotImplementedError(f"{handler_type} handlers is not yet supported.")

    constraints = [None] * len(topologies) if constraints is None else constraints

    converter = _CONVERTERS[handler_type]
    converter_spec = inspect.signature(converter.fn)
    converter_kwargs = {}

    if "topologies" in converter_spec.parameters:
        converter_kwargs["topologies"] = topologies
    if "v_site_maps" in converter_spec.parameters:
        assert v_site_maps is not None, "v-site maps must be provided"
        converter_kwargs["v_site_maps"] = v_site_maps
    if "constraints" in converter_spec.parameters:
        constraint_idxs = [[] if v is None else v.idxs.tolist() for v in constraints]
        unique_idxs = [{tuple(sorted(idxs)) for idxs in v} for v in constraint_idxs]

        converter_kwargs["constraints"] = unique_idxs

    potentials_by_type = (
        {}
        if potentials is None
        else {potential.type: (potential, maps) for potential, maps in potentials}
    )

    dependencies = {}
    depends_on = converter.depends_on if converter.depends_on is not None else []

    if len(depends_on) > 0:
        missing_deps = {dep for dep in depends_on if dep not in potentials_by_type}
        assert len(missing_deps) == 0, "missing dependencies"

        dependencies = {dep: potentials_by_type[dep] for dep in depends_on}
        assert "dependencies" in converter_spec.parameters, "dependencies not accepted"

    if "dependencies" in converter_spec.parameters:
        converter_kwargs["dependencies"] = dependencies

    converted = converter.fn(handlers, **converter_kwargs)
    converted = [converted] if not isinstance(converted, list) else converted

    converted_by_type = {
        potential.type: (potential, maps) for potential, maps in converted
    }
    assert len(converted_by_type) == len(converted), "duplicate potentials found"

    potentials_by_type = {
        **{
            potential.type: (potential, maps)
            for potential, maps in potentials_by_type.values()
            if potential.type not in depends_on
            and potential.type not in converted_by_type
        },
        **converted_by_type,
    }

    return [*potentials_by_type.values()]

convert_interchange

convert_interchange(
    interchange: Interchange | list[Interchange],
) -> tuple[TensorForceField, list[TensorTopology]]

Convert a list of interchange objects into tensor potentials.

Parameters:

• interchange (Interchange | list[Interchange]) –

  The list of (or singile) interchange objects to convert into tensor potentials.

Returns:

• tuple[TensorForceField, list[TensorTopology]] –

  The tensor force field containing the parameters of each handler, and a list (one per interchange) of objects mapping molecule elements (e.g. bonds, angles) to corresponding handler parameters.

Examples:

>>> from openff.toolkit import ForceField, Molecule
>>> from openff.interchange import Interchange
>>>
>>> force_field = ForceField("openff_unconstrained-2.0.0.offxml")
>>> molecules = [Molecule.from_smiles("CCO"), Molecule.from_smiles("CC")]
>>>
>>> interchanges = [
...     Interchange.from_smirnoff(force_field, molecule.to_topology())
...     for molecule in molecules
... ]
>>>
>>> tensor_ff, tensor_topologies = convert_interchange(interchanges)

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

def convert_interchange(
    interchange: openff.interchange.Interchange | list[openff.interchange.Interchange],
) -> tuple[smee.TensorForceField, list[smee.TensorTopology]]:
    """Convert a list of interchange objects into tensor potentials.

    Args:
        interchange: The list of (or singile) interchange objects to convert into
            tensor potentials.

    Returns:
        The tensor force field containing the parameters of each handler, and a list
        (one per interchange) of objects mapping molecule elements (e.g. bonds, angles)
        to corresponding handler parameters.

    Examples:

        >>> from openff.toolkit import ForceField, Molecule
        >>> from openff.interchange import Interchange
        >>>
        >>> force_field = ForceField("openff_unconstrained-2.0.0.offxml")
        >>> molecules = [Molecule.from_smiles("CCO"), Molecule.from_smiles("CC")]
        >>>
        >>> interchanges = [
        ...     Interchange.from_smirnoff(force_field, molecule.to_topology())
        ...     for molecule in molecules
        ... ]
        >>>
        >>> tensor_ff, tensor_topologies = convert_interchange(interchanges)
    """
    importlib.import_module("smee.converters.openff.nonbonded")
    importlib.import_module("smee.converters.openff.valence")

    interchanges = (
        [interchange]
        if isinstance(interchange, openff.interchange.Interchange)
        else interchange
    )
    topologies = []

    handler_types = {
        handler_type
        for interchange in interchanges
        for handler_type in interchange.collections
    }
    handlers_by_type = {handler_type: [] for handler_type in sorted(handler_types)}

    for interchange in interchanges:
        for handler_type in handlers_by_type:
            handler = (
                None
                if handler_type not in interchange.collections
                else interchange.collections[handler_type]
            )
            handlers_by_type[handler_type].append(handler)

        topologies.append(interchange.topology)

    v_sites, v_site_maps = None, [None] * len(topologies)

    if "VirtualSites" in handlers_by_type:
        v_sites, v_site_maps = _convert_v_sites(
            handlers_by_type.pop("VirtualSites"), topologies
        )

    constraints = [None] * len(topologies)

    if "Constraints" in handlers_by_type:
        constraints = _convert_constraints(handlers_by_type.pop("Constraints"))

    conversion_order = _resolve_conversion_order([*handlers_by_type])
    converted = []

    for handler_type in conversion_order:
        handlers = handlers_by_type[handler_type]

        if (
            sum(len(handler.potentials) for handler in handlers if handler is not None)
            == 0
        ):
            continue

        converted = convert_handlers(
            handlers, topologies, v_site_maps, converted, constraints
        )

    # handlers may either return multiple potentials, or condense multiple already
    # converted potentials into a single one (e.g. electrostatics into some polarizable
    # potential)
    potentials = []
    parameter_maps_by_handler = {}

    for potential, parameter_maps in converted:
        potentials.append(potential)
        parameter_maps_by_handler[potential.type] = parameter_maps

    tensor_topologies = [
        _convert_topology(
            topology,
            {
                potential.type: parameter_maps_by_handler[potential.type][i]
                for potential in potentials
            },
            v_site_maps[i],
            constraints[i],
        )
        for i, topology in enumerate(topologies)
    ]

    tensor_force_field = smee.TensorForceField(potentials, v_sites)
    return tensor_force_field, tensor_topologies

smirnoff_parameter_converter

smirnoff_parameter_converter(
    type_: str,
    default_units: dict[str, Unit],
    depends_on: list[str] | None = None,
)

A decorator used to flag a function as being able to convert a parameter handlers parameters into tensors.

Parameters:

• type_ (str) –

  The type of parameter handler that the decorated function can convert.

• default_units (dict[str, Unit]) –

  The default units of each parameter in the handler.

• depends_on (list[str] | None, default: None ) –

  The names of other handlers that this handler depends on. When set, the convert function should additionally take in a list of the already converted potentials and return a new list of potentials that should either include or replace the original potentials.

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

def smirnoff_parameter_converter(
    type_: str,
    default_units: dict[str, openff.units.Unit],
    depends_on: list[str] | None = None,
):
    """A decorator used to flag a function as being able to convert a parameter handlers
    parameters into tensors.

    Args:
        type_: The type of parameter handler that the decorated function can convert.
        default_units: The default units of each parameter in the handler.
        depends_on: The names of other handlers that this handler depends on. When set,
            the convert function should additionally take in a list of the already
            converted potentials and return a new list of potentials that should either
            include or replace the original potentials.
    """

    def parameter_converter_inner(func):
        if type_ in _CONVERTERS:
            raise KeyError(f"A {type_} converter is already registered.")

        _CONVERTERS[type_] = _Converter(func, default_units, depends_on)

        return func

    return parameter_converter_inner