Don't hardcode built-in attrs, and fix docs

Makefile

@@ -48,4 +48,4 @@ endif
 
 docs-insiders:
 	$(CONDA_ENV_RUN) pip install git+https://$(INSIDER_DOCS_TOKEN)@github.com/SimonBoothroyd/mkdocstrings-python.git \
-                    			 git+https://$(INSIDER_DOCS_TOKEN)@github.com/SimonBoothroyd/griffe-pydantic.git@fix-inheritence-static
+                    			 git+https://$(INSIDER_DOCS_TOKEN)@github.com/SimonBoothroyd/griffe-pydantic.git

mkdocs.yml

@@ -24,7 +24,6 @@ nav:
   - examples/README.md
   - examples/compute-energy.ipynb
   - examples/conformer-minimization.ipynb
-  - examples/custom-handler.ipynb
   - examples/md-simulations.ipynb
   - examples/parameter-gradients.ipynb
 - API reference: reference/
@@ -117,7 +116,6 @@ plugins:
           show_signature_annotations: true
           show_symbol_type_heading: true
           show_symbol_type_toc: true
-          show_if_no_docstring: false
           signature_crossrefs: true
           summary: true
           members_order: source

smee/__init__.py

@@ -5,7 +5,7 @@
 """
 
 from . import _version
-from ._constants import EnergyFn, PotentialType
+from ._constants import CUTOFF_ATTRIBUTE, SWITCH_ATTRIBUTE, EnergyFn, PotentialType
 from ._models import (
     NonbondedParameterMap,
     ParameterMap,
@@ -24,6 +24,8 @@
 __version__ = _version.get_versions()["version"]
 
 __all__ = [
+    "CUTOFF_ATTRIBUTE",
+    "SWITCH_ATTRIBUTE",
     "EnergyFn",
     "PotentialType",
     "ValenceParameterMap",

smee/_constants.py

@@ -54,3 +54,13 @@ class EnergyFn(_StrEnum):
     ANGLE_HARMONIC = "k/2*(theta-angle)**2"
 
     TORSION_COSINE = "k*(1+cos(periodicity*theta-phase))"
+
+
+CUTOFF_ATTRIBUTE = "cutoff"
+"""The attribute that should be used to store the cutoff distance of a potential."""
+SWITCH_ATTRIBUTE = "switch_width"
+"""The attribute that should be used to store the switch width of a potential, if the
+potential should use the standard OpenMM switch function.
+
+This attribute should be omitted if the potential should not use a switch function.
+"""

smee/converters/openmm/nonbonded.py

@@ -38,11 +38,11 @@ def _create_nonbonded_force(
     else:
         raise NotImplementedError(f"unsupported force class {cls}")
 
-    cutoff_idx = potential.attribute_cols.index("cutoff")
+    cutoff_idx = potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)
     switch_idx = (
         None
-        if "switch_width" not in potential.attribute_cols
-        else potential.attribute_cols.index("switch_width")
+        if smee.SWITCH_ATTRIBUTE not in potential.attribute_cols
+        else potential.attribute_cols.index(smee.SWITCH_ATTRIBUTE)
     )
 
     if not system.is_periodic:
@@ -505,9 +505,9 @@ def convert_dexp_potential(
     The intermolcular interactions are described by a custom nonbonded force, while the
     intramolecular interactions are described by a custom bond force.
 
-    If the potential has custom mixing rules (i.e. exceptions), the a lookup table will
-    be used to store the parameters. Otherwise, the mixing rules will be applied
-    directly in the energy function.
+    If the potential has custom mixing rules (i.e. exceptions), a lookup table will be
+    used to store the parameters. Otherwise, the mixing rules will be applied directly
+    in the energy function.
     """
     energy_fn = (
         "epsilon * (repulsion - attraction);"

smee/converters/openmm/valence.py

@@ -16,6 +16,7 @@
 def convert_bond_potential(
     potential: smee.TensorPotential, system: smee.TensorSystem
 ) -> openmm.HarmonicBondForce:
+    """Convert a harmonic bond potential to a corresponding OpenMM force."""
     force = openmm.HarmonicBondForce()
 
     idx_offset = 0
@@ -47,6 +48,7 @@ def convert_bond_potential(
 def _convert_angle_potential(
     potential: smee.TensorPotential, system: smee.TensorSystem
 ) -> openmm.HarmonicAngleForce:
+    """Convert a harmonic angle potential to a corresponding OpenMM force."""
     force = openmm.HarmonicAngleForce()
 
     idx_offset = 0
@@ -82,6 +84,7 @@ def _convert_angle_potential(
 def convert_torsion_potential(
     potential: smee.TensorPotential, system: smee.TensorSystem
 ) -> openmm.PeriodicTorsionForce:
+    """Convert a torsion potential to a corresponding OpenMM force."""
     force = openmm.PeriodicTorsionForce()
 
     idx_offset = 0

smee/potentials/_potentials.py

@@ -264,8 +264,10 @@ def _precompute_pairwise(
 
         requires_pairwise = True
 
-        if "cutoff" in potential.attribute_cols:
-            cutoff = potential.attributes[potential.attribute_cols.index("cutoff")]
+        if smee.CUTOFF_ATTRIBUTE in potential.attribute_cols:
+            cutoff = potential.attributes[
+                potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)
+            ]
             cutoffs.append(cutoff)
 
         break

smee/potentials/nonbonded.py

@@ -270,10 +270,10 @@ def _compute_switch_fn(
     potential: smee.TensorPotential,
     pairwise: PairwiseDistances,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
-    if "switch_width" not in potential.attribute_cols:
+    if smee.SWITCH_ATTRIBUTE not in potential.attribute_cols:
         return torch.ones(1), None
 
-    switch_width_idx = potential.attribute_cols.index("switch_width")
+    switch_width_idx = potential.attribute_cols.index(smee.SWITCH_ATTRIBUTE)
     switch_width = pairwise.cutoff - potential.attributes[switch_width_idx]
 
     x_switch = (pairwise.distances - switch_width) / (pairwise.cutoff - switch_width)
@@ -454,7 +454,7 @@ def compute_lj_energy(
 
     box_vectors = None if not system.is_periodic else box_vectors
 
-    cutoff = potential.attributes[potential.attribute_cols.index("cutoff")]
+    cutoff = potential.attributes[potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)]
 
     pairwise = (
         pairwise
@@ -728,7 +728,7 @@ def compute_dexp_energy(
     """
     box_vectors = None if not system.is_periodic else box_vectors
 
-    cutoff = potential.attributes[potential.attribute_cols.index("cutoff")]
+    cutoff = potential.attributes[potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)]
 
     pairwise = (
         pairwise
@@ -906,7 +906,7 @@ def _compute_coulomb_energy_periodic(
 
     charges = smee.potentials.broadcast_parameters(system, potential).squeeze(-1)
 
-    cutoff = potential.attributes[potential.attribute_cols.index("cutoff")]
+    cutoff = potential.attributes[potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)]
     error_tol = torch.tensor(0.0001)
 
     exceptions = _compute_pme_exclusions(system, potential).to(charges.device)
@@ -994,7 +994,7 @@ def compute_coulomb_energy(
 
     box_vectors = None if not system.is_periodic else box_vectors
 
-    cutoff = potential.attributes[potential.attribute_cols.index("cutoff")]
+    cutoff = potential.attributes[potential.attribute_cols.index(smee.CUTOFF_ATTRIBUTE)]
 
     pairwise = (
         pairwise

smee/tests/convertors/openff/test_nonbonded.py

@@ -32,7 +32,7 @@ def test_convert_electrostatics_am1bcc(ethanol, ethanol_interchange):
         "scale_13",
         "scale_14",
         "scale_15",
-        "cutoff",
+        smee.CUTOFF_ATTRIBUTE,
     )
 
     assert potential.parameter_cols == ("charge",)