converters

Convert to / from smee tensor representations.

Modules:

• openff –

  Tensor representations of SMIRNOFF force fields.

• openmm –

  Convert tensor representations into OpenMM systems.

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_to_openmm_force –

  Convert a smee potential to OpenMM forces.

• convert_to_openmm_system –

  Convert a smee force field and system / topology into an OpenMM system.

• convert_to_openmm_topology –

  Convert a smee system to an OpenMM topology.

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,
) -> 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.

• 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,
) -> 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.
        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.")

    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

    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)

    # 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

convert_to_openmm_force

convert_to_openmm_force(
    potential: TensorPotential, system: TensorSystem
) -> list[Force]

Convert a smee potential to OpenMM forces.

Some potentials may return multiple forces, e.g. a vdW potential may return one force containing intermolecular interactions and another containing intramolecular interactions.

See Also

potential_converter: for how to define a converter function.

Parameters:

• potential (TensorPotential) –

  The potential to convert.

• system (TensorSystem) –

  The system to convert.

Returns:

• list[Force] –

  The OpenMM force(s).

Source code in smee/converters/openmm/_openmm.py

def convert_to_openmm_force(
    potential: smee.TensorPotential, system: smee.TensorSystem
) -> list[openmm.Force]:
    """Convert a ``smee`` potential to OpenMM forces.

    Some potentials may return multiple forces, e.g. a vdW potential may return one
    force containing intermolecular interactions and another containing intramolecular
    interactions.

    See Also:
        potential_converter: for how to define a converter function.

    Args:
        potential: The potential to convert.
        system: The system to convert.

    Returns:
        The OpenMM force(s).
    """
    # register the built-in converter functions
    importlib.import_module("smee.converters.openmm.nonbonded")
    importlib.import_module("smee.converters.openmm.valence")

    potential = potential.to("cpu")
    system = system.to("cpu")

    if potential.exceptions is not None and potential.type != "vdW":
        raise NotImplementedError("exceptions are only supported for vdW potentials")

    converter_key = (str(potential.type), str(potential.fn))

    if converter_key not in _CONVERTER_FUNCTIONS:
        raise NotImplementedError(
            f"cannot convert type={potential.type} fn={potential.fn} to an OpenMM force"
        )

    forces = _CONVERTER_FUNCTIONS[converter_key](potential, system)
    return forces if isinstance(forces, (list, tuple)) else [forces]

convert_to_openmm_system

convert_to_openmm_system(
    force_field: TensorForceField,
    system: TensorSystem | TensorTopology,
) -> System

Convert a smee force field and system / topology into an OpenMM system.

Parameters:

• force_field (TensorForceField) –

  The force field parameters.

• system (TensorSystem | TensorTopology) –

  The system / topology to convert.

Returns:

• System –

  The OpenMM system.

Source code in smee/converters/openmm/_openmm.py

def convert_to_openmm_system(
    force_field: smee.TensorForceField,
    system: smee.TensorSystem | smee.TensorTopology,
) -> openmm.System:
    """Convert a ``smee`` force field and system / topology into an OpenMM system.

    Args:
        force_field: The force field parameters.
        system: The system / topology to convert.

    Returns:
        The OpenMM system.
    """

    system: smee.TensorSystem = (
        system
        if isinstance(system, smee.TensorSystem)
        else smee.TensorSystem([system], [1], False)
    )

    force_field = force_field.to("cpu")
    system = system.to("cpu")

    omm_forces = {
        potential_type: convert_to_openmm_force(potential, system)
        for potential_type, potential in force_field.potentials_by_type.items()
    }
    omm_system = create_openmm_system(system, force_field.v_sites)

    if (
        "Electrostatics" in omm_forces
        and "vdW" in omm_forces
        and len(omm_forces["vdW"]) == 1
        and isinstance(omm_forces["vdW"][0], openmm.NonbondedForce)
    ):
        (electrostatic_force,) = omm_forces.pop("Electrostatics")
        (vdw_force,) = omm_forces.pop("vdW")

        nonbonded_force = _combine_nonbonded(vdw_force, electrostatic_force)
        omm_system.addForce(nonbonded_force)

    for forces in omm_forces.values():
        for force in forces:
            omm_system.addForce(force)

    _apply_constraints(omm_system, system)

    return omm_system

convert_to_openmm_topology

convert_to_openmm_topology(
    system: TensorSystem | TensorTopology,
) -> Topology

Convert a smee system to an OpenMM topology.

Notes

Virtual sites are given the element symbol "X" and atomic number 82.

Parameters:

• system (TensorSystem | TensorTopology) –

  The system to convert.

Returns:

• Topology –

  The OpenMM topology.

Source code in smee/converters/openmm/_openmm.py

def convert_to_openmm_topology(
    system: smee.TensorSystem | smee.TensorTopology,
) -> openmm.app.Topology:
    """Convert a ``smee`` system to an OpenMM topology.

    Notes:
        Virtual sites are given the element symbol "X" and atomic number 82.

    Args:
        system: The system to convert.

    Returns:
        The OpenMM topology.
    """
    system: smee.TensorSystem = (
        system
        if isinstance(system, smee.TensorSystem)
        else smee.TensorSystem([system], [1], False)
    )

    omm_topology = openmm.app.Topology()

    for topology, n_copies in zip(system.topologies, system.n_copies, strict=True):
        chain = omm_topology.addChain()

        is_water = topology.n_atoms == 3 and sorted(
            int(v) for v in topology.atomic_nums
        ) == [1, 1, 8]

        residue_name = "WAT" if is_water else "UNK"

        for _ in range(n_copies):
            residue = omm_topology.addResidue(residue_name, chain)
            element_counter = collections.defaultdict(int)

            atoms = {}

            for i, atomic_num in enumerate(topology.atomic_nums):
                element = openmm.app.Element.getByAtomicNumber(int(atomic_num))
                element_counter[element.symbol] += 1

                name = element.symbol + (
                    ""
                    if element_counter[element.symbol] == 1 and element.symbol != "H"
                    else f"{element_counter[element.symbol]}"
                )
                atoms[i] = omm_topology.addAtom(name, element, residue)

            for _ in range(topology.n_v_sites):
                omm_topology.addAtom(
                    "X", openmm.app.Element.getByAtomicNumber(82), residue
                )

            for bond_idxs, bond_order in zip(
                topology.bond_idxs, topology.bond_orders, strict=True
            ):
                idx_a, idx_b = int(bond_idxs[0]), int(bond_idxs[1])

                bond_order = int(bond_order)
                bond_type = {
                    1: openmm.app.Single,
                    2: openmm.app.Double,
                    3: openmm.app.Triple,
                }[bond_order]

                omm_topology.addBond(atoms[idx_a], atoms[idx_b], bond_type, bond_order)

    return omm_topology