convert

Convert OpenFF ForceField <--> smee TensorForceField.

Functions:

• convert_to_smirnoff –

  Convert a tensor force field that contains bespoke valence parameters to

• parameterise –

  Prepare a Trainable object that contains a force field with

• linearise_harmonics_force_field –

  Linearize the harmonic potential parameters in the forcefield.

• linearise_harmonics_topology –

  Linearize harmonic potential parameters in the topology.

_reflect_angle

_reflect_angle(angle: float) -> float

Reflect an angle (in radians) to be in the range [0, pi).

Source code in presto/convert.py

def _reflect_angle(angle: float) -> float:
    """Reflect an angle (in radians) to be in the range [0, pi)."""
    return math.pi - abs((angle % (2 * math.pi)) - math.pi)

convert_to_smirnoff

convert_to_smirnoff(
    ff: TensorForceField, base: ForceField | None = None
) -> ForceField

Convert a tensor force field that contains bespoke valence parameters to SMIRNOFF format. Args: ff: The force field containing the bespoke valence terms. base: The (optional) original SMIRNOFF force field to add the bespoke parameters to. If no specified, a force field containing only the bespoke parameters will be returned. Returns: A SMIRNOFF force field containing the valence terms of the input force field.

Source code in presto/convert.py

def convert_to_smirnoff(
    ff: smee.TensorForceField, base: openff.toolkit.ForceField | None = None
) -> openff.toolkit.ForceField:
    """Convert a tensor force field that *contains bespoke valence parameters* to
    SMIRNOFF format.
    Args:
        ff: The force field containing the bespoke valence terms.
        base: The (optional) original SMIRNOFF force field to add the bespoke
            parameters to. If no specified, a force field containing only the bespoke
            parameters will be returned.
    Returns:
        A SMIRNOFF force field containing the valence terms of the input force field.
    """
    ff_smirnoff = openff.toolkit.ForceField() if base is None else copy.deepcopy(base)

    for potential in ff.potentials:
        if potential.type in {"Bonds", "Angles", "ProperTorsions", "ImproperTorsions"}:
            assert potential.attribute_cols is None
            parameters_by_smarts: dict[str, dict[int | None, torch.Tensor]] = (
                collections.defaultdict(dict)
            )
            for parameter, parameter_key in zip(
                potential.parameters, potential.parameter_keys, strict=True
            ):
                assert parameter_key.mult not in parameters_by_smarts[parameter_key.id]
                parameters_by_smarts[parameter_key.id][parameter_key.mult] = parameter
            handler = ff_smirnoff.get_parameter_handler(potential.type)
            for smarts, parameters_by_mult in parameters_by_smarts.items():
                mults = {*parameters_by_mult}
                if None in mults and len(mults) > 1:
                    raise NotImplementedError("unexpected parameters found")
                if None not in mults and mults != {*range(len(mults))}:
                    raise NotImplementedError("unexpected parameters found")
                counter = len(handler.parameters) + 1
                parameter_id = f"{potential.type[0].lower()}-bespoke-{counter}"
                parameter_dict: dict[str, str | Quantity] = {
                    "smirks": smarts,
                    "id": parameter_id,
                }
                parameter_dict.update(
                    {
                        (col if mult is None else f"{col}{mult + 1}"): float(
                            parameter[col_idx]
                        )
                        * potential.parameter_units[col_idx]
                        for mult, parameter in parameters_by_mult.items()
                        for col_idx, col in enumerate(potential.parameter_cols)
                    }
                )
                _add_parameter_with_overwrite(handler, parameter_dict)

        elif potential.type == "LinearBonds":
            assert potential.attribute_cols is None
            parameters_by_smarts = collections.defaultdict(dict)
            new_params = []
            for param in potential.parameters:
                k1 = param[0].item()
                k2 = param[1].item()
                b1 = param[2].item()
                b2 = param[3].item()
                k = k1 + k2
                b = (k1 * b1 + k2 * b2) / k
                dt = param.dtype
                new_params.append([k, b])
            reconstructed_param = torch.tensor(new_params, dtype=dt)
            reconstructed_units = (_KCAL_PER_MOL_ANGSQ, _ANGSTROM)
            reconstructed_cols = ("k", "length")
            for parameter, parameter_key in zip(
                reconstructed_param, potential.parameter_keys, strict=True
            ):
                assert parameter_key.mult not in parameters_by_smarts[parameter_key.id]
                parameters_by_smarts[parameter_key.id][parameter_key.mult] = parameter
            handler = ff_smirnoff.get_parameter_handler("Bonds")
            for smarts, parameters_by_mult in parameters_by_smarts.items():
                mults = {*parameters_by_mult}
                if None in mults and len(mults) > 1:
                    raise NotImplementedError("unexpected parameters found")
                if None not in mults and mults != {*range(len(mults))}:
                    raise NotImplementedError("unexpected parameters found")
                counter = len(handler.parameters) + 1
                parameter_id = f"{potential.type[0].lower()}-bespoke-{counter}"
                parameter_dict = {"smirks": smarts, "id": parameter_id}
                parameter_dict.update(
                    {
                        (col if mult is None else f"{col}{mult + 1}"): float(
                            parameter[col_idx]
                        )
                        * reconstructed_units[col_idx]
                        for mult, parameter in parameters_by_mult.items()
                        for col_idx, col in enumerate(reconstructed_cols)
                    }
                )
                _add_parameter_with_overwrite(handler, parameter_dict)
        elif potential.type == "LinearAngles":
            assert potential.attribute_cols is None
            parameters_by_smarts = collections.defaultdict(dict)
            new_params = []
            for param in potential.parameters:
                k1 = param[0].item()
                k2 = param[1].item()
                a1 = param[2].item()
                a2 = param[3].item()
                k = k1 + k2
                # Set k and angle to 0 if very close
                a = (k1 * a1 + k2 * a2) / k
                # Ensure that the angle is in the range [0, pi)
                a = _reflect_angle(a)
                dt = param.dtype
                new_params.append([k, a])
            reconstructed_param = torch.tensor(new_params, dtype=dt)
            reconstructed_units = (_KCAL_PER_MOL_RADSQ, _RADIANS)
            reconstructed_cols = ("k", "angle")
            for parameter, parameter_key in zip(
                reconstructed_param, potential.parameter_keys, strict=True
            ):
                assert parameter_key.mult not in parameters_by_smarts[parameter_key.id]
                parameters_by_smarts[parameter_key.id][parameter_key.mult] = parameter
            handler = ff_smirnoff.get_parameter_handler("Angles")
            for smarts, parameters_by_mult in parameters_by_smarts.items():
                mults = {*parameters_by_mult}
                if None in mults and len(mults) > 1:
                    raise NotImplementedError("unexpected parameters found")
                if None not in mults and mults != {*range(len(mults))}:
                    raise NotImplementedError("unexpected parameters found")
                counter = len(handler.parameters) + 1
                parameter_id = f"{potential.type[0].lower()}-bespoke-{counter}"
                parameter_dict = {"smirks": smarts, "id": parameter_id}
                parameter_dict.update(
                    {
                        (col if mult is None else f"{col}{mult + 1}"): float(
                            parameter[col_idx]
                        )
                        * reconstructed_units[col_idx]
                        for mult, parameter in parameters_by_mult.items()
                        for col_idx, col in enumerate(reconstructed_cols)
                    }
                )
                _add_parameter_with_overwrite(handler, parameter_dict)

    return ff_smirnoff

parameterise

parameterise(
    settings: ParameterisationSettings,
    device: TorchDevice = "cuda",
) -> tuple[
    list[Molecule],
    ForceField,
    list[TensorTopology],
    TensorForceField,
]

Prepare a Trainable object that contains a force field with unique parameters for each topologically symmetric term across multiple molecules.

Parameters:

• settings (ParameterisationSettings) –

  The settings for the parameterisation.

• device (TorchDevice, default: 'cuda' ) –

  The device to use for the force field and topology.

Returns:

• mols ( list[Molecule] ) –

  The molecules that have been parameterised.

• off_ff ( ForceField ) –

  The original force field, used as a base for the bespoke force field.

• tensor_tops ( list[TensorTopology] ) –

  The topologies of the molecules.

• tensor_ff ( TensorForceField ) –

  The force field with unique parameters for each topologically symmetric term.

Source code in presto/convert.py

def parameterise(
    settings: ParameterisationSettings,
    device: TorchDevice = "cuda",
) -> tuple[
    list[openff.toolkit.Molecule],
    openff.toolkit.ForceField,
    list[smee.TensorTopology],
    smee.TensorForceField,
]:
    """Prepare a Trainable object that contains a force field with
    unique parameters for each topologically symmetric term across multiple molecules.

    Parameters
    ----------
    settings: ParameterisationSettings
        The settings for the parameterisation.

    device: TorchDevice, default "cuda"
        The device to use for the force field and topology.

    Returns
    -------
    mols: list[openff.toolkit.Molecule]
        The molecules that have been parameterised.
    off_ff: openff.toolkit.ForceField
        The original force field, used as a base for the bespoke force field.
    tensor_tops: list[smee.TensorTopology]
        The topologies of the molecules.
    tensor_ff: smee.TensorForceField
        The force field with unique parameters for each topologically
        symmetric term.
    """
    # Create molecules from SMILES
    mols = settings.molecules

    off_ff = openff.toolkit.ForceField(settings.initial_force_field)

    if "[#1:1]-[*:2]" in off_ff["Constraints"].parameters:
        logger.warning(
            "The force field contains a constraint for [#1:1]-[*:2] which "
            "is not supported. Removing this constraint."
        )
        del off_ff["Constraints"].parameters["[#1:1]-[*:2]"]

    if settings.expand_torsions:
        torsion_generation_settings = settings.type_generation_settings.get(
            "ProperTorsions"
        )
        excluded_smirks = (
            torsion_generation_settings.exclude
            if torsion_generation_settings is not None
            else None
        )
        off_ff = _expand_torsions(off_ff, excluded_smirks=excluded_smirks)

    # Add bespoke parameters to the force field (deduplicated across all molecules)
    bespoke_ff = add_types_to_forcefield(
        mols,
        off_ff,
        settings.type_generation_settings,
    )

    if settings.msm_settings is not None:
        bespoke_ff = apply_msm_to_molecules(
            mols=mols,
            off_ff=bespoke_ff,
            settings=settings.msm_settings,
        )

    # Create separate Interchange objects for each molecule
    interchanges = [
        openff.interchange.Interchange.from_smirnoff(bespoke_ff, mol.to_topology())
        for mol in mols
    ]

    # Convert all interchanges at once to get a shared force field
    # and separate topologies that all reference the same parameters
    force_field, tensor_tops = smee.converters.convert_interchange(interchanges)
    force_field = force_field.to(device)

    # Move each topology to the device
    tensor_tops = [top.to(device) for top in tensor_tops]

    if settings.linearise_harmonics:
        force_field = linearise_harmonics_force_field(force_field, device)
        for i in range(len(tensor_tops)):
            tensor_tops[i] = linearise_harmonics_topology(tensor_tops[i], device)

    return (
        mols,
        bespoke_ff,
        tensor_tops,
        force_field,
    )

_expand_torsions

_expand_torsions(
    ff: ForceField, excluded_smirks: list[str] | None = None
) -> ForceField

Expand the torsion potential to include K0-4 for proper torsions

Source code in presto/convert.py

def _expand_torsions(
    ff: openff.toolkit.ForceField, excluded_smirks: list[str] | None = None
) -> openff.toolkit.ForceField:
    """Expand the torsion potential to include K0-4 for proper torsions"""
    excluded_smirks = excluded_smirks or []
    ff_copy = copy.deepcopy(ff)
    torsion_handler = ff_copy.get_parameter_handler("ProperTorsions")
    for parameter in torsion_handler:
        # Avoid expanding any excluded smarts
        if parameter.smirks in excluded_smirks:
            continue
        # set the defaults
        parameter.idivf = [1.0] * 4
        default_k = [0 * _KCAL_PER_MOL] * 4
        default_phase = [0 * _RADIANS] * 4
        default_p = [1, 2, 3, 4]
        # update the existing k values for the correct phase and p
        for i, p in enumerate(parameter.periodicity):
            try:
                default_k[p - 1] = parameter.k[i]
                default_phase[p - 1] = parameter.phase[i]
            except IndexError:
                continue
        # update with new parameters
        parameter.k = default_k
        parameter.phase = default_phase
        parameter.periodicity = default_p
    return ff_copy

_add_angle_within_range

_add_angle_within_range(
    initial_angle: float, diff: float
) -> float

Add a difference to an angle cap to be within [0, pi]

Source code in presto/convert.py

def _add_angle_within_range(initial_angle: float, diff: float) -> float:
    """Add a difference to an angle cap to be within [0, pi]"""
    new_angle = initial_angle + diff
    if diff > 0:
        return min(new_angle, math.pi)
    else:
        return max(new_angle, 0.0)

_compute_linear_harmonic_params

_compute_linear_harmonic_params(
    k: float,
    eq_value: float,
    compute_lower_bound: Callable[[float], float],
    compute_upper_bound: Callable[[float], float],
) -> tuple[float, float, float, float]

Compute linearized harmonic parameters from standard parameters.

This generic function distributes a force constant across two bounds, inversely proportional to the distance from each bound.

Args: k: Force constant (e.g., kcal/mol/Ų or kcal/mol/rad²) eq_value: Equilibrium value (e.g., bond length or angle) compute_lower_bound: Function that takes eq_value and returns lower bound compute_upper_bound: Function that takes eq_value and returns upper bound

Returns: Tuple of (k1, k2, eq1, eq2) where: - k1, k2: Distributed force constants - eq1, eq2: Lower and upper equilibrium value bounds

Source code in presto/convert.py

def _compute_linear_harmonic_params(
    k: float,
    eq_value: float,
    compute_lower_bound: Callable[[float], float],
    compute_upper_bound: Callable[[float], float],
) -> tuple[float, float, float, float]:
    """Compute linearized harmonic parameters from standard parameters.

    This generic function distributes a force constant across two bounds,
    inversely proportional to the distance from each bound.

    Args:
        k: Force constant (e.g., kcal/mol/Ų or kcal/mol/rad²)
        eq_value: Equilibrium value (e.g., bond length or angle)
        compute_lower_bound: Function that takes eq_value and returns
            lower bound
        compute_upper_bound: Function that takes eq_value and returns
            upper bound

    Returns:
        Tuple of (k1, k2, eq1, eq2) where:
        - k1, k2: Distributed force constants
        - eq1, eq2: Lower and upper equilibrium value bounds
    """
    eq1 = compute_lower_bound(eq_value)
    eq2 = compute_upper_bound(eq_value)
    d = eq2 - eq1
    # Distribute force constant inversely proportional to distance from bounds
    k1 = k * (eq2 - eq_value) / d
    k2 = k * (eq_value - eq1) / d
    return k1, k2, eq1, eq2

_linearize_bond_parameters

_linearize_bond_parameters(
    potential: TensorPotential, device_type: str
) -> TensorPotential

Linearize bond potential parameters.

Converts standard harmonic bond parameters (k, length) to linearized form (k1, k2, b1, b2) where the equilibrium bond length range is [0.5length, 1.5length].

Source code in presto/convert.py

def _linearize_bond_parameters(
    potential: smee.TensorPotential, device_type: str
) -> smee.TensorPotential:
    """Linearize bond potential parameters.

    Converts standard harmonic bond parameters (k, length) to linearized
    form (k1, k2, b1, b2) where the equilibrium bond length range is
    [0.5*length, 1.5*length].
    """
    new_potential = copy.deepcopy(potential)
    new_potential.type = "LinearBonds"
    new_potential.fn = "(k1+k2)/2*(r-(k1*length1+k2*length2)/(k1+k2))**2"
    new_potential.parameter_cols = ("k1", "k2", "b1", "b2")

    # Get dtype from the first parameter
    dtype = potential.parameters.dtype

    new_params = [
        _compute_linear_harmonic_params(
            param[0].item(),
            param[1].item(),
            lambda b: b - 0.4,  # Lower bound: current length - 0.4 Å
            lambda b: b + 0.4,  # Upper bound: current length + 0.4 Å
        )
        for param in potential.parameters
    ]

    new_potential.parameters = torch.tensor(
        new_params, dtype=dtype, requires_grad=False, device=device_type
    )
    new_potential.parameter_units = (
        _KCAL_PER_MOL_ANGSQ,
        _KCAL_PER_MOL_ANGSQ,
        _ANGSTROM,
        _ANGSTROM,
    )
    return new_potential

_linearize_angle_parameters

_linearize_angle_parameters(
    potential: TensorPotential, device_type: str
) -> TensorPotential

Linearize angle potential parameters.

Converts standard harmonic angle parameters (k, angle) to linearized form (k1, k2, angle1, angle2) where the equilibrium angle range is [0, π].

Source code in presto/convert.py

def _linearize_angle_parameters(
    potential: smee.TensorPotential, device_type: str
) -> smee.TensorPotential:
    """Linearize angle potential parameters.

    Converts standard harmonic angle parameters (k, angle) to linearized form
    (k1, k2, angle1, angle2) where the equilibrium angle range is [0, π].
    """
    new_potential = copy.deepcopy(potential)
    new_potential.type = "LinearAngles"
    new_potential.fn = "(k1+k2)/2*(r-(k1*angle1+k2*angle2)/(k1+k2))**2"
    new_potential.parameter_cols = ("k1", "k2", "angle1", "angle2")

    # Get dtype from the first parameter
    dtype = potential.parameters.dtype

    new_params = [
        _compute_linear_harmonic_params(
            param[0].item(),
            param[1].item(),
            lambda a: max(0.0, a - math.pi / 3),  # Lower bound: max(0, angle - π/3)
            lambda a: min(math.pi, a + math.pi / 3),  # Upper bound: min(π, angle + π/3)
        )
        for param in potential.parameters
    ]

    new_potential.parameters = torch.tensor(
        new_params, dtype=dtype, requires_grad=False, device=device_type
    )
    new_potential.parameter_units = (
        _KCAL_PER_MOL_RADSQ,
        _KCAL_PER_MOL_RADSQ,
        _RADIANS,
        _RADIANS,
    )
    return new_potential

linearise_harmonics_force_field

linearise_harmonics_force_field(
    ff: TensorForceField, device_type: str
) -> TensorForceField

Linearize the harmonic potential parameters in the forcefield.

This converts Bonds and Angles potentials to their linearized forms (LinearBonds and LinearAngles) for more robust optimization.

Source code in presto/convert.py

def linearise_harmonics_force_field(
    ff: smee.TensorForceField, device_type: str
) -> smee.TensorForceField:
    """Linearize the harmonic potential parameters in the forcefield.

    This converts Bonds and Angles potentials to their linearized forms
    (LinearBonds and LinearAngles) for more robust optimization.
    """
    ff_copy = copy.deepcopy(ff)
    ff_copy.potentials = []

    for potential in ff.potentials:
        if potential.type == "Bonds":
            ff_copy.potentials.append(
                _linearize_bond_parameters(potential, device_type)
            )
        elif potential.type == "Angles":
            ff_copy.potentials.append(
                _linearize_angle_parameters(potential, device_type)
            )
        else:
            ff_copy.potentials.append(potential)

    return ff_copy

linearise_harmonics_topology

linearise_harmonics_topology(
    topology: TensorTopology, device_type: TorchDevice
) -> TensorTopology

Linearize harmonic potential parameters in the topology.

This updates the topology to use LinearBonds and LinearAngles parameter maps instead of Bonds and Angles.

Source code in presto/convert.py

def linearise_harmonics_topology(
    topology: smee.TensorTopology, device_type: TorchDevice
) -> smee.TensorTopology:
    """Linearize harmonic potential parameters in the topology.

    This updates the topology to use LinearBonds and LinearAngles
    parameter maps instead of Bonds and Angles.
    """
    topology_copy = topology.to(device_type)
    if "Bonds" in topology_copy.parameters:
        topology_copy.parameters["LinearBonds"] = copy.deepcopy(
            topology_copy.parameters["Bonds"]
        )
        del topology_copy.parameters["Bonds"]
    if "Angles" in topology_copy.parameters:
        topology_copy.parameters["LinearAngles"] = copy.deepcopy(
            topology_copy.parameters["Angles"]
        )
        del topology_copy.parameters["Angles"]
    return topology_copy