[WIP] Add DampedExp6810 and multipole handlers / potentials

.gitignore

@@ -113,3 +113,4 @@ ENV/
 
 # Local development
 scratch
+CLAUDE.md

smee/_constants.py

@@ -40,6 +40,7 @@ class EnergyFn(_StrEnum):
     """An enumeration of the energy functions supported by ``smee`` out of the box."""
 
     COULOMB = "coul"
+    POLARIZATION = "coul+pol"
 
     VDW_LJ = "4*epsilon*((sigma/r)**12-(sigma/r)**6)"
     VDW_DEXP = (
@@ -48,6 +49,7 @@ class EnergyFn(_StrEnum):
         "alpha/(alpha-beta)*exp(beta*(1-r/r_min)))"
     )
     # VDW_BUCKINGHAM = "a*exp(-b*r)-c*r^-6"
+    VDW_DAMPEDEXP6810 = "force_at_zero*beta**-1*exp(-beta*(r-rho))-f_6(beta*r)*c6**6-f_8(beta*r)*c8**8-f_10(beta*r)*c10**10"
 
     BOND_HARMONIC = "k/2*(r-length)**2"
 

smee/converters/openff/_openff.py

@@ -71,6 +71,26 @@ def _get_value(
 ) -> float:
     """Returns the value of a parameter in its default units"""
     default_units = default_units[parameter]
+
+    # Handle missing parameters by using default values
+    if parameter not in potential.parameters:
+        if default_value is None:
+            # Set specific defaults for known multipole parameters
+            if parameter == "thole":
+                default_value = 0.39 * openff.units.unit.dimensionless
+            elif parameter == "dampingFactor":
+                # Will be computed from polarizability if not provided
+                return 0.0  # Placeholder, will be computed later
+            elif parameter.startswith("dipole"):
+                default_value = 0.0 * openff.units.unit.elementary_charge * openff.units.unit.angstrom
+            elif parameter.startswith("quadrupole"):
+                default_value = 0.0 * openff.units.unit.elementary_charge * openff.units.unit.angstrom**2
+            elif parameter in ["axisType", "multipoleAtomZ", "multipoleAtomX", "multipoleAtomY"]:
+                default_value = -1 * openff.units.unit.dimensionless  # -1 indicates undefined
+            else:
+                raise KeyError(f"Parameter '{parameter}' not found and no default provided")
+        return default_value.m_as(default_units)
+
     value = potential.parameters[parameter]
     return (value if value is not None else default_value).m_as(default_units)
 

smee/converters/openff/nonbonded.py

@@ -217,6 +217,282 @@ def convert_dexp(
     return potential, parameter_maps
 
 
+@smee.converters.smirnoff_parameter_converter(
+    "DampedExp6810",
+    {
+        "rho": _ANGSTROM,
+        "beta": _ANGSTROM**-1,
+        "c6": _KCAL_PER_MOL * _ANGSTROM**6,
+        "c8": _KCAL_PER_MOL * _ANGSTROM**8,
+        "c10": _KCAL_PER_MOL * _ANGSTROM**10,
+        "force_at_zero": _KCAL_PER_MOL * _ANGSTROM**-1,
+        "scale_12": _UNITLESS,
+        "scale_13": _UNITLESS,
+        "scale_14": _UNITLESS,
+        "scale_15": _UNITLESS,
+        "cutoff": _ANGSTROM,
+        "switch_width": _ANGSTROM,
+    },
+)
+def convert_dampedexp6810(
+    handlers: list[
+        "smirnoff_plugins.collections.nonbonded.SMIRNOFFDampedExp6810Collection"
+    ],
+    topologies: list[openff.toolkit.Topology],
+    v_site_maps: list[smee.VSiteMap | None],
+) -> tuple[smee.TensorPotential, list[smee.NonbondedParameterMap]]:
+    import smee.potentials.nonbonded
+
+    (
+        potential,
+        parameter_maps,
+    ) = smee.converters.openff.nonbonded.convert_nonbonded_handlers(
+        handlers,
+        "DampedExp6810",
+        topologies,
+        v_site_maps,
+        ("rho", "beta", "c6", "c8", "c10"),
+        ("cutoff", "switch_width", "force_at_zero"),
+    )
+    potential.type = smee.PotentialType.VDW
+    potential.fn = smee.EnergyFn.VDW_DAMPEDEXP6810
+
+    return potential, parameter_maps
+
+
+@smee.converters.smirnoff_parameter_converter(
+    "Multipole",
+    {
+        # Molecular multipole moments
+        "dipoleX": _ELEMENTARY_CHARGE * _ANGSTROM,
+        "dipoleY": _ELEMENTARY_CHARGE * _ANGSTROM,
+        "dipoleZ": _ELEMENTARY_CHARGE * _ANGSTROM,
+        "quadrupoleXX": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        "quadrupoleXY": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        "quadrupoleXZ": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        "quadrupoleYX": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        "quadrupoleYY": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        "quadrupoleYZ": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        "quadrupoleZX": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        "quadrupoleZY": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        "quadrupoleZZ": _ELEMENTARY_CHARGE * _ANGSTROM**2,
+        # Local frame definition
+        "axisType": _UNITLESS,
+        "multipoleAtomZ": _UNITLESS,
+        "multipoleAtomX": _UNITLESS,
+        "multipoleAtomY": _UNITLESS,
+        # Damping and polarizability (these may not be present in current force fields)
+        "thole": _UNITLESS,
+        "dampingFactor": _ANGSTROM,
+        "polarity": _ANGSTROM**3,
+        # Cutoff and scaling
+        "cutoff": _ANGSTROM,
+        "scale_12": _UNITLESS,
+        "scale_13": _UNITLESS,
+        "scale_14": _UNITLESS,
+        "scale_15": _UNITLESS,
+    },
+    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,
+        (
+            "dipoleX", "dipoleY", "dipoleZ",
+            "quadrupoleXX", "quadrupoleXY", "quadrupoleXZ",
+            "quadrupoleYX", "quadrupoleYY", "quadrupoleYZ",
+            "quadrupoleZX", "quadrupoleZY", "quadrupoleZZ",
+            "axisType", "multipoleAtomZ", "multipoleAtomX", "multipoleAtomY",
+            "thole", "dampingFactor", "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,
+    ]
+
+    # Handle different numbers of columns between charge and polarizability potentials
+    n_chg_cols = potential_chg.parameters.shape[1]
+    n_pol_cols = potential_pol.parameters.shape[1]
+
+    # Pad charge parameters with zeros for the new polarizability columns
+    parameters_chg = torch.cat(
+        (potential_chg.parameters, torch.zeros(potential_chg.parameters.shape[0], n_pol_cols, dtype=potential_chg.parameters.dtype)), dim=1
+    )
+
+    # Resolve multipole axis atoms from SMIRKS indices to actual topology indices.
+    # The OFFXML stores multipoleAtomZ/X/Y as 1-based SMIRKS atom map numbers.
+    # We need to resolve these to actual 0-based topology indices for each atom.
+    # Columns in potential_pol.parameters: 13=multipoleAtomZ, 14=multipoleAtomX, 15=multipoleAtomY
+
+    parameter_key_to_idx = {
+        key: i for i, key in enumerate(potential_pol.parameter_keys)
+    }
+
+    all_resolved_pol_params = []
+    resolved_parameter_maps_pol = []
+
+    for handler, topology, v_site_map, param_map_pol in zip(
+        handlers, topologies, v_site_maps, parameter_maps_pol, strict=True
+    ):
+        n_atoms = topology.n_atoms
+
+        # Create per-atom resolved parameters for multipoles
+        # Each atom gets its own parameter row with resolved axis indices
+        resolved_params = []
+
+        for atom_idx in range(n_atoms):
+            # Find the topology_key for this atom
+            matched_key = None
+            matched_param_idx = None
+
+            for topology_key, parameter_key in handler.key_map.items():
+                if isinstance(topology_key, openff.interchange.models.VirtualSiteKey):
+                    continue
+                if topology_key.atom_indices[0] == atom_idx:
+                    matched_key = topology_key
+                    matched_param_idx = parameter_key_to_idx[parameter_key]
+                    break
+
+            if matched_key is None:
+                # No multipole parameters for this atom, create zero row
+                resolved_params.append(torch.zeros(n_pol_cols, dtype=torch.float64))
+                continue
+
+            # Get the base parameters for this atom
+            base_params = potential_pol.parameters[matched_param_idx].clone()
+
+            # Resolve SMIRKS indices to actual topology indices
+            # SMIRKS indices are 1-based, so multipoleAtomZ=2 means atom_indices[1]
+            smirks_atom_z = int(base_params[13])  # multipoleAtomZ
+            smirks_atom_x = int(base_params[14])  # multipoleAtomX
+            smirks_atom_y = int(base_params[15])  # multipoleAtomY
+
+            # Resolve using atom_indices from the match
+            # SMIRKS :N corresponds to atom_indices[N-1]
+            atom_indices = matched_key.atom_indices
+
+            if smirks_atom_z > 0 and smirks_atom_z <= len(atom_indices):
+                base_params[13] = atom_indices[smirks_atom_z - 1]
+            else:
+                base_params[13] = -1
+
+            if smirks_atom_x > 0 and smirks_atom_x <= len(atom_indices):
+                base_params[14] = atom_indices[smirks_atom_x - 1]
+            else:
+                base_params[14] = -1
+
+            if smirks_atom_y > 0 and smirks_atom_y <= len(atom_indices):
+                base_params[15] = atom_indices[smirks_atom_y - 1]
+            else:
+                base_params[15] = -1
+
+            # Compute damping factor from polarity
+            base_params[17] = base_params[18] ** (1/6)
+
+            resolved_params.append(base_params)
+
+        # Stack into a tensor for this topology
+        resolved_params_tensor = torch.stack(resolved_params)
+        all_resolved_pol_params.append(resolved_params_tensor)
+
+        # Create identity-like assignment matrix (each atom maps to its own row)
+        assignment_matrix = torch.eye(n_atoms, dtype=torch.float64).to_sparse()
+        resolved_parameter_maps_pol.append(
+            smee.NonbondedParameterMap(
+                assignment_matrix=assignment_matrix,
+                exclusions=param_map_pol.exclusions,
+                exclusion_scale_idxs=param_map_pol.exclusion_scale_idxs,
+            )
+        )
+
+    # Combine all resolved parameters across topologies into a single tensor
+    # Each topology's parameters are separate, so we need to track offsets
+    total_pol_params = sum(p.shape[0] for p in all_resolved_pol_params)
+
+    # Create combined parameters tensor
+    if all_resolved_pol_params:
+        combined_pol_params = torch.cat(all_resolved_pol_params, dim=0)
+    else:
+        combined_pol_params = torch.zeros((0, n_pol_cols), dtype=torch.float64)
+
+    # Pad with charge columns
+    parameters_pol = torch.cat(
+        (torch.zeros(combined_pol_params.shape[0], n_chg_cols, dtype=torch.float64), combined_pol_params), dim=1
+    )
+
+    potential_chg.parameters = torch.cat((parameters_chg, parameters_pol), dim=0)
+
+    # Update assignment matrices with proper offsets
+    param_offset = 0
+    for i, (parameter_map_chg, resolved_map_pol) in enumerate(zip(
+        parameter_maps_chg, resolved_parameter_maps_pol, strict=True
+    )):
+        n_chg_params = parameter_map_chg.assignment_matrix.shape[1]
+        n_atoms = resolved_map_pol.assignment_matrix.shape[0]
+
+        # The resolved_map_pol assignment matrix is identity, but we need to offset
+        # the column indices by (n_charge_params + param_offset)
+        pol_assignment_dense = resolved_map_pol.assignment_matrix.to_dense()
+
+        # Create the full assignment matrix
+        n_total_params = parameters_chg.shape[0] + combined_pol_params.shape[0]
+        full_assignment = torch.zeros((n_atoms, n_total_params), dtype=torch.float64)
+
+        # Copy charge assignments
+        chg_assignment = parameter_map_chg.assignment_matrix.to_dense()
+        full_assignment[:, :n_chg_params] = chg_assignment
+
+        # Add multipole assignments with proper offset
+        pol_start_col = parameters_chg.shape[0] + param_offset
+        for atom_idx in range(n_atoms):
+            full_assignment[atom_idx, pol_start_col + atom_idx] = 1.0
+
+        parameter_map_chg.assignment_matrix = full_assignment.to_sparse()
+        param_offset += n_atoms
+
+    return potential_chg, parameter_maps_chg
+
+
 def _make_v_site_electrostatics_compatible(
     handlers: list[openff.interchange.smirnoff.SMIRNOFFElectrostaticsCollection],
 ):