[WIP] Add DampedExp6810 and multipole handlers / potentials

[aehogan]

Description

Support for DampedExp6810 and Multipole from the smirnoff-plugins repository.

Status

• Ready to go

[aehogan]

@SimonBoothroyd To get the explicit polarization potential working I'm going to need the electrostatic potential (dV/dr)*q from the Coulomb potential. Right now potential_energy_fn only takes as inputs the TensorSystem and its own TensorPotential.

Is there an easy way I can get access to the electrostatics potential? Perhaps just the charges from that potential?

The way this is structured in a .offxml / smirnoff plugins is to have both an Electrostatics and Multipole tag, when constructing an openmm system the multipole handler is loaded second, copies the charges from the nonbonded handler to the AmoebaMultipoleForce and then zeros out the nonbonded charges.

[SimonBoothroyd]

ok @aehogan you'll need to rebase onto main, but everything should pretty much be there now! i think something like this should work for the converter:

@smee.converters.smirnoff_parameter_converter(
    "Multipole",
    {"polarity": _ANGSTROM**3, "cutoff": _ANGSTROM},
    depends_on=["Electrostatics"],
)
def convert_multipole(
    handlers: list[
        "smirnoff_plugins.collections.nonbonded.SMIRNOFFMultipoleCollection"
    ],
    topologies: list[openff.toolkit.Topology],
    v_site_maps: list[smee.VSiteMap | None],
    dependencies: dict[
        str, tuple[smee.TensorPotential, list[smee.NonbondedParameterMap]]
    ],
) -> tuple[smee.TensorPotential, list[smee.NonbondedParameterMap]]:

    potential_chg, parameter_maps_chg = dependencies["Electrostatics"]

    (
        potential_pol,
        parameter_maps_pol,
    ) = smee.converters.openff.nonbonded.convert_nonbonded_handlers(
        handlers,
        "Multipole",
        topologies,
        v_site_maps,
        ("polarity",),
        ("cutoff",),
        has_exclusions=False,
    )

    cutoff_idx_pol = potential_pol.attribute_cols.index("cutoff")
    cutoff_idx_chg = potential_chg.attribute_cols.index("cutoff")

    assert torch.isclose(
        potential_pol.attributes[cutoff_idx_pol],
        potential_chg.attributes[cutoff_idx_chg],
    )

    potential_chg.fn = smee.EnergyFn.POLARIZATION

    potential_chg.parameter_cols = (
        *potential_chg.parameter_cols,
        *potential_pol.parameter_cols,
    )
    potential_chg.parameter_units = (
        *potential_chg.parameter_units,
        *potential_pol.parameter_units,
    )
    potential_chg.parameter_keys = [
        *potential_chg.parameter_keys,
        *potential_pol.parameter_keys,
    ]

    parameters_chg = torch.cat(
        (potential_chg.parameters, torch.zeros_like(potential_chg.parameters)), dim=1
    )
    parameters_pol = torch.cat(
        (torch.zeros_like(potential_pol.parameters), potential_pol.parameters), dim=1
    )
    potential_chg.parameters = torch.cat((parameters_chg, parameters_pol), dim=0)

    for parameter_map_chg, parameter_map_pol in zip(
        parameter_maps_chg, parameter_maps_pol, strict=True
    ):
        parameter_map_chg.assignment_matrix = torch.block_diag(
            parameter_map_chg.assignment_matrix.to_dense(),
            parameter_map_pol.assignment_matrix.to_dense(),
        ).to_sparse()

    return potential_chg, parameter_maps_chg