diff --git a/pyproject.toml b/pyproject.toml
index b38c0977..5afcf8f0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -219,21 +219,3 @@ warn_unused_configs = true
 disallow_untyped_defs = false  # too strict for existing codebase
 check_untyped_defs = true
 ignore_missing_imports = true
-
-# Grandfathered modules: mypy is enforced repo-wide via the CI Lint job, but
-# the modules listed below carry pre-existing type errors. They are still
-# followed for type info, but their own errors are not reported — so any NEW
-# module, or any module removed from this list, must be mypy-clean.
-# Un-grandfather modules one at a time as their errors are fixed.
-# See https://mypy.readthedocs.io/en/stable/existing_code.html
-[[tool.mypy.overrides]]
-ignore_errors = true
-module = [
-  "q2mm.backends.mm",
-  "q2mm.backends.mm._jax_common",
-  "q2mm.backends.mm.batched",
-  "q2mm.backends.mm.jax_engine",
-  "q2mm.backends.mm.jax_md_engine",
-  "q2mm.backends.mm.openmm",
-  "q2mm.backends.mm.tinker",
-]
diff --git a/q2mm/backends/conformance.py b/q2mm/backends/conformance.py
index bc697b91..031b073e 100644
--- a/q2mm/backends/conformance.py
+++ b/q2mm/backends/conformance.py
@@ -411,6 +411,51 @@ def _validate_result(
 }
 
 
+def _verify_undeclared_batched_hessian(
+    descriptor: BackendDescriptor,
+    backend: Backend,
+    prepared: object,
+) -> Capability:
+    try:
+        prepare_hessian_batches(backend, [prepared])  # type: ignore[list-item]
+    except UnsupportedCapabilityError:
+        return Capability.BATCHED_HESSIAN
+    except Exception as exc:  # noqa: BLE001
+        raise ConformanceError(
+            f"{descriptor.name}: undeclared capability batched_hessian raised "
+            f"{type(exc).__name__}, expected UnsupportedCapabilityError."
+        ) from exc
+    raise ConformanceError(
+        f"{descriptor.name}: undeclared capability batched_hessian did not raise UnsupportedCapabilityError."
+    )
+
+
+def _verify_reusable_state(
+    descriptor: BackendDescriptor,
+    reusable: tuple[Callable[[object], object], object, type, Capability] | None,
+    *,
+    n_atoms: int,
+    n_params: int,
+    kind: str,
+) -> Capability:
+    if reusable is None:
+        raise ConformanceError(
+            f"{descriptor.name}: reusable_state was selected but no selected {kind} capability can prove reuse."
+        )
+    method, request, result_type, capability = reusable
+    for _ in range(2):
+        result = _call_declared(descriptor, capability, method, request)
+        _validate_result(
+            descriptor,
+            capability,
+            result,
+            result_type,
+            n_atoms=n_atoms,
+            n_params=n_params,
+        )
+    return Capability.REUSABLE_STATE
+
+
 def run_mm_conformance(case: MMConformanceCase) -> ConformanceOutcome:
     """Run the selected API-v1 MM checks and all undeclared-operation gates."""
     if not isinstance(case, MMConformanceCase):
@@ -480,40 +525,18 @@ def run_mm_conformance(case: MMConformanceCase) -> ConformanceOutcome:
                 ) from exc
             executed.append(Capability.BATCHED_HESSIAN)
     else:
-        try:
-            prepare_hessian_batches(case.backend, [prepared])
-        except UnsupportedCapabilityError:
-            unsupported.append(Capability.BATCHED_HESSIAN)
-        except Exception as exc:  # noqa: BLE001
-            raise ConformanceError(
-                f"{case.descriptor.name}: undeclared capability batched_hessian raised "
-                f"{type(exc).__name__}, expected UnsupportedCapabilityError."
-            ) from exc
-        else:
-            raise ConformanceError(
-                f"{case.descriptor.name}: undeclared capability batched_hessian did not raise "
-                "UnsupportedCapabilityError."
-            )
+        unsupported.append(_verify_undeclared_batched_hessian(case.descriptor, case.backend, prepared))
 
     if Capability.REUSABLE_STATE in info.capabilities:
-        if reusable is None:
-            raise ConformanceError(
-                f"{case.descriptor.name}: reusable_state was selected but no selected prepared-session "
-                "capability can prove reuse."
-            )
-        method, request, result_type, capability = reusable
-        first = _call_declared(case.descriptor, capability, method, request)
-        second = _call_declared(case.descriptor, capability, method, request)
-        for result in (first, second):
-            _validate_result(
+        executed.append(
+            _verify_reusable_state(
                 case.descriptor,
-                capability,
-                result,
-                result_type,
+                reusable,
                 n_atoms=case.molecule.n_atoms,
                 n_params=len(vector),
+                kind="prepared-session",
             )
-        executed.append(Capability.REUSABLE_STATE)
+        )
 
     _assert_unmutated(case.descriptor.name, case.molecule, molecule_before)
     _assert_unmutated(case.descriptor.name, case.force_field, force_field_before)
@@ -560,39 +583,18 @@ def run_reference_conformance(case: ReferenceConformanceCase) -> ConformanceOutc
             _prove_unsupported(case.descriptor, capability, method, request)
             unsupported.append(capability)
 
-    try:
-        prepare_hessian_batches(case.backend, [prepared])
-    except UnsupportedCapabilityError:
-        unsupported.append(Capability.BATCHED_HESSIAN)
-    except Exception as exc:  # noqa: BLE001
-        raise ConformanceError(
-            f"{case.descriptor.name}: undeclared capability batched_hessian raised "
-            f"{type(exc).__name__}, expected UnsupportedCapabilityError."
-        ) from exc
-    else:
-        raise ConformanceError(
-            f"{case.descriptor.name}: undeclared capability batched_hessian did not raise UnsupportedCapabilityError."
-        )
+    unsupported.append(_verify_undeclared_batched_hessian(case.descriptor, case.backend, prepared))
 
     if Capability.REUSABLE_STATE in info.capabilities:
-        if reusable is None:
-            raise ConformanceError(
-                f"{case.descriptor.name}: reusable_state was selected but no selected reference capability can prove reuse."
-            )
-        method, request, result_type, capability = reusable
-        for result in (
-            _call_declared(case.descriptor, capability, method, request),
-            _call_declared(case.descriptor, capability, method, request),
-        ):
-            _validate_result(
+        executed.append(
+            _verify_reusable_state(
                 case.descriptor,
-                capability,
-                result,
-                result_type,
+                reusable,
                 n_atoms=case.molecule.n_atoms,
                 n_params=0,
+                kind="reference",
             )
-        executed.append(Capability.REUSABLE_STATE)
+        )
 
     _assert_unmutated(case.descriptor.name, case.molecule, molecule_before)
     return ConformanceOutcome(
diff --git a/q2mm/backends/contracts.py b/q2mm/backends/contracts.py
index 3eaabcc3..61016b24 100644
--- a/q2mm/backends/contracts.py
+++ b/q2mm/backends/contracts.py
@@ -51,7 +51,7 @@
 from collections.abc import Collection, Mapping, Sequence
 from dataclasses import dataclass, field
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Protocol, runtime_checkable
+from typing import TYPE_CHECKING, Literal, Protocol, runtime_checkable
 
 import numpy as np
 import numpy.typing as npt
@@ -139,6 +139,30 @@ def _role_capability_conflicts(
     return frozenset(capabilities) - _ROLE_CAPABILITIES[role]
 
 
+def _normalize_capabilities_and_forms(
+    role: BackendRole,
+    capabilities: Collection[Capability],
+    functional_forms: Collection[str],
+    *,
+    owner: str,
+    verb: str,
+) -> tuple[frozenset[Capability], frozenset[str]]:
+    """Validate and freeze capability and functional-form sets for *role*."""
+    caps = frozenset(capabilities)
+    if not all(isinstance(c, Capability) for c in caps):
+        raise ValueError(f"{owner} capabilities must contain only Capability members.")
+    conflicts = _role_capability_conflicts(role, caps)
+    if conflicts:
+        values = sorted(capability.value for capability in conflicts)
+        raise ValueError(f"{owner}: role {role.value!r} cannot {verb} capabilities {values}.")
+    forms = frozenset(functional_forms)
+    if not all(isinstance(f, str) for f in forms):
+        raise ValueError(f"{owner} functional forms must contain only strings.")
+    if role is BackendRole.REFERENCE and forms:
+        raise ValueError(f"{owner}: reference backends must declare no functional forms.")
+    return caps, forms
+
+
 class EnergyUnit(str, enum.Enum):
     """Explicit canonical energy unit for a result."""
 
@@ -240,21 +264,15 @@ def __post_init__(self) -> None:
             raise ValueError("BackendInfo.name must be a non-empty string.")
         if not isinstance(self.role, BackendRole):
             raise ValueError(f"BackendInfo.role must be a BackendRole; got {self.role!r}.")
-        # Normalize capabilities/forms to frozensets of the right element types.
-        caps = frozenset(self.capabilities)
-        if not all(isinstance(c, Capability) for c in caps):
-            raise ValueError("BackendInfo.capabilities must all be Capability members.")
-        conflicts = _role_capability_conflicts(self.role, caps)
-        if conflicts:
-            values = sorted(capability.value for capability in conflicts)
-            raise ValueError(f"BackendInfo: role {self.role.value!r} cannot declare capabilities {values}.")
-        forms = frozenset(self.functional_forms)
-        if not all(isinstance(f, str) for f in forms):
-            raise ValueError("BackendInfo.functional_forms must all be strings.")
+        caps, forms = _normalize_capabilities_and_forms(
+            self.role,
+            self.capabilities,
+            self.functional_forms,
+            owner="BackendInfo",
+            verb="declare",
+        )
         object.__setattr__(self, "capabilities", caps)
         object.__setattr__(self, "functional_forms", forms)
-        if self.role is BackendRole.REFERENCE and forms:
-            raise ValueError("BackendInfo: reference backends must declare no functional_forms.")
         if self.provenance is not None:
             if not isinstance(self.provenance, BackendProvenance):
                 raise ValueError("BackendInfo.provenance must be a BackendProvenance or None.")
@@ -467,8 +485,8 @@ def __post_init__(self) -> None:
 
 
 @dataclass(frozen=True, eq=False)
-class EnergyRequest:
-    """Single-point energy for a full parameter vector."""
+class _ParameterVectorRequest:
+    """Base for MM requests carrying one validated full parameter vector."""
 
     parameters: np.ndarray
 
@@ -477,7 +495,12 @@ def __post_init__(self) -> None:
 
 
 @dataclass(frozen=True, eq=False)
-class MinimizationRequest:
+class EnergyRequest(_ParameterVectorRequest):
+    """Single-point energy for a full parameter vector."""
+
+
+@dataclass(frozen=True, eq=False)
+class MinimizationRequest(_ParameterVectorRequest):
     """Energy-minimize (relax) the geometry for a full parameter vector.
 
     Args:
@@ -489,28 +512,22 @@ class MinimizationRequest:
 
     """
 
-    parameters: np.ndarray
     max_iterations: int | None = None
     tolerance: float | None = None
 
     def __post_init__(self) -> None:
-        object.__setattr__(self, "parameters", _readonly_vector(self.parameters, name="parameters"))
+        super().__post_init__()
         if self.max_iterations is not None and self.max_iterations <= 0:
             raise EvaluationError("MinimizationRequest.max_iterations must be positive.")
 
 
 @dataclass(frozen=True, eq=False)
-class HessianRequest:
+class HessianRequest(_ParameterVectorRequest):
     """Cartesian Hessian for a full parameter vector."""
 
-    parameters: np.ndarray
-
-    def __post_init__(self) -> None:
-        object.__setattr__(self, "parameters", _readonly_vector(self.parameters, name="parameters"))
-
 
 @dataclass(frozen=True, eq=False)
-class FrequencyRequest:
+class FrequencyRequest(_ParameterVectorRequest):
     """Vibrational frequencies for a full parameter vector.
 
     Args:
@@ -520,32 +537,18 @@ class FrequencyRequest:
 
     """
 
-    parameters: np.ndarray
     on_error: str = "raise"
 
-    def __post_init__(self) -> None:
-        object.__setattr__(self, "parameters", _readonly_vector(self.parameters, name="parameters"))
-
 
 @dataclass(frozen=True, eq=False)
-class ParameterGradientRequest:
+class ParameterGradientRequest(_ParameterVectorRequest):
     """Energy plus analytical ``dE/dp`` for a full parameter vector."""
 
-    parameters: np.ndarray
-
-    def __post_init__(self) -> None:
-        object.__setattr__(self, "parameters", _readonly_vector(self.parameters, name="parameters"))
-
 
 @dataclass(frozen=True, eq=False)
-class HessianJacobianRequest:
+class HessianJacobianRequest(_ParameterVectorRequest):
     """Hessian plus its analytical ``dH/dp`` Jacobian for a full vector."""
 
-    parameters: np.ndarray
-
-    def __post_init__(self) -> None:
-        object.__setattr__(self, "parameters", _readonly_vector(self.parameters, name="parameters"))
-
 
 @dataclass(frozen=True, eq=False)
 class BatchedEnergyRequest:
@@ -563,7 +566,7 @@ def __post_init__(self) -> None:
 
 
 @dataclass(frozen=True, eq=False)
-class BatchedHessianRequest:
+class BatchedHessianRequest(_ParameterVectorRequest):
     """Cartesian Hessians for a batch of topology-compatible prepared cases.
 
     Carries one full parameter vector applied to every case in the batch; the
@@ -575,11 +578,6 @@ class BatchedHessianRequest:
 
     """
 
-    parameters: np.ndarray
-
-    def __post_init__(self) -> None:
-        object.__setattr__(self, "parameters", _readonly_vector(self.parameters, name="parameters"))
-
 
 # ---------------------------------------------------------------------------
 # Reference evaluation requests (no parameter vectors — model fixed at prepare)
@@ -1125,6 +1123,41 @@ def layout(self) -> ParameterLayout:
             raise UnsupportedCapabilityError(self._info.name, Capability.ENERGY)
         return self._layout
 
+    @property
+    def _provenance(self) -> BackendProvenance:
+        """Return this session's non-null BackendProvenance."""
+        if self._info.provenance is None:
+            raise EvaluationError(f"{self._info.name}: session info has no BackendProvenance.")
+        return self._info.provenance
+
+    def _ff_for(self, parameters: np.ndarray) -> ForceField:
+        """Validate *parameters* against :attr:`layout` and return an updated :class:`ForceField`."""
+        vec = self._validate_vector(parameters)
+        return self.layout.replace(self.force_field, vec)
+
+    def _raw_hessian(self, ff: ForceField) -> np.ndarray:
+        """Compute the raw atomic-unit Cartesian Hessian for *ff* (MM subclasses)."""
+        raise UnsupportedCapabilityError(self._info.name, Capability.HESSIAN)
+
+    def _mm_frequencies(self, request: FrequencyRequest, *, backend_label: str) -> FrequencyResult:
+        """Compute MM vibrational frequencies by converting :meth:`_raw_hessian` via :func:`hessian_to_frequencies`."""
+        from q2mm.models.hessian import hessian_to_frequencies
+
+        ff = self._ff_for(request.parameters)
+        try:
+            hess_au = self._raw_hessian(ff)
+            on_error: Literal["raise", "penalty"] = "penalty" if request.on_error == "penalty" else "raise"
+            if request.on_error not in ("raise", "penalty"):
+                raise ValueError(f"Invalid on_error mode: {request.on_error!r}")
+            freqs = hessian_to_frequencies(hess_au, list(self.molecule.symbols), on_error=on_error)
+        except Exception as exc:  # noqa: BLE001
+            raise EvaluationError(f"{backend_label} frequency evaluation failed: {exc}") from exc
+        return FrequencyResult(
+            frequencies=readonly_array(freqs),
+            unit=FrequencyUnit.INVERSE_CM,
+            provenance=self._provenance,
+        )
+
     # -- validation helpers -------------------------------------------------
 
     def _require(self, capability: Capability) -> None:
@@ -1523,20 +1556,13 @@ def __post_init__(self) -> None:
             raise ValueError("BackendDescriptor.name must be non-empty.")
         if not isinstance(self.role, BackendRole):
             raise ValueError("BackendDescriptor.role must be a BackendRole.")
-        capabilities = frozenset(self.capability_ceiling)
-        if not all(isinstance(capability, Capability) for capability in capabilities):
-            raise ValueError("BackendDescriptor.capability_ceiling must contain only Capability members.")
-        forms = frozenset(self.functional_form_ceiling)
-        if not all(isinstance(form, str) for form in forms):
-            raise ValueError("BackendDescriptor.functional_form_ceiling must contain only strings.")
-        if self.role is BackendRole.REFERENCE and forms:
-            raise ValueError("BackendDescriptor reference backends must have an empty functional_form_ceiling.")
-        conflicts = _role_capability_conflicts(self.role, capabilities)
-        if conflicts:
-            values = sorted(capability.value for capability in conflicts)
-            raise ValueError(
-                f"BackendDescriptor role {self.role.value!r} cannot include capabilities {values} in its ceiling."
-            )
+        capabilities, forms = _normalize_capabilities_and_forms(
+            self.role,
+            self.capability_ceiling,
+            self.functional_form_ceiling,
+            owner="BackendDescriptor",
+            verb="include",
+        )
         object.__setattr__(self, "capability_ceiling", capabilities)
         object.__setattr__(self, "functional_form_ceiling", forms)
         if not isinstance(self.factory, str) or ":" not in self.factory:
diff --git a/q2mm/backends/discovery.py b/q2mm/backends/discovery.py
index 2c186a37..11c21baa 100644
--- a/q2mm/backends/discovery.py
+++ b/q2mm/backends/discovery.py
@@ -291,6 +291,14 @@ def __post_init__(self) -> None:
             raise ValueError("ManifestFailure.message must be a non-empty string.")
 
 
+def _normalize_records(records: Iterable[DiscoveryRecord], *, owner: str) -> tuple[DiscoveryRecord, ...]:
+    """Validate and sort discovery records into a deterministic tuple."""
+    items = tuple(records)
+    if not all(isinstance(record, DiscoveryRecord) for record in items):
+        raise ValueError(f"{owner}.records must all be DiscoveryRecord instances.")
+    return tuple(sorted(items, key=lambda record: record.sort_key))
+
+
 @dataclass(frozen=True, eq=False)
 class DiscoverySnapshot:
     """Immutable result of one discovery pass.
@@ -313,11 +321,8 @@ def __post_init__(self) -> None:
             if key != descriptor.name:
                 raise ValueError(f"DiscoverySnapshot key {key!r} must equal descriptor name {descriptor.name!r}.")
             frozen[key] = descriptor
-        records = tuple(self.records)
-        if not all(isinstance(record, DiscoveryRecord) for record in records):
-            raise ValueError("DiscoverySnapshot.records must all be DiscoveryRecord instances.")
         object.__setattr__(self, "descriptors", MappingProxyType(frozen))
-        object.__setattr__(self, "records", tuple(sorted(records, key=lambda record: record.sort_key)))
+        object.__setattr__(self, "records", _normalize_records(self.records, owner="DiscoverySnapshot"))
 
 
 @dataclass(frozen=True)
@@ -334,10 +339,7 @@ class DiscoveryReport:
 
     def __post_init__(self) -> None:
         """Validate and normalize records to a deterministically-ordered tuple."""
-        records = tuple(self.records)
-        if not all(isinstance(record, DiscoveryRecord) for record in records):
-            raise ValueError("DiscoveryReport.records must all be DiscoveryRecord instances.")
-        object.__setattr__(self, "records", tuple(sorted(records, key=lambda record: record.sort_key)))
+        object.__setattr__(self, "records", _normalize_records(self.records, owner="DiscoveryReport"))
 
     @property
     def registered(self) -> tuple[str, ...]:
@@ -377,9 +379,7 @@ def _as_str_sequence(value: object) -> list[str] | None:
 
 def _is_json_safe(value: object) -> bool:
     """Return whether *value* uses the closed JSON data vocabulary."""
-    if value is None or isinstance(value, (str, bool)):
-        return True
-    if isinstance(value, int):
+    if value is None or isinstance(value, (str, int)):
         return True
     if isinstance(value, float):
         return math.isfinite(value)
diff --git a/q2mm/backends/mm/__init__.py b/q2mm/backends/mm/__init__.py
index f3f7a792..4dbe594d 100644
--- a/q2mm/backends/mm/__init__.py
+++ b/q2mm/backends/mm/__init__.py
@@ -14,7 +14,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from q2mm.backends.mm.openmm import OpenMMBackend
 from q2mm.backends.mm.tinker import TinkerBackend
@@ -31,7 +31,7 @@
 }
 
 
-def __getattr__(name: str) -> type:
+def __getattr__(name: str) -> Any:
     if name in _LAZY_IMPORTS:
         import importlib
 
diff --git a/q2mm/backends/mm/_jax_common.py b/q2mm/backends/mm/_jax_common.py
index 9c11fedc..469085f8 100644
--- a/q2mm/backends/mm/_jax_common.py
+++ b/q2mm/backends/mm/_jax_common.py
@@ -27,11 +27,36 @@
 
 import importlib.util
 from collections.abc import Sequence
+from dataclasses import dataclass
 from types import ModuleType
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+import numpy as np
 
 from q2mm._jax_support import has_jax, load_jax
+from q2mm.backends.contracts import (
+    AbstractPreparedBackend,
+    BackendInfo,
+    BatchedEnergyRequest,
+    BatchedEnergyResult,
+    EnergyRequest,
+    EnergyResult,
+    EnergyUnit,
+    EvaluationError,
+    FrequencyRequest,
+    FrequencyResult,
+    HessianRequest,
+    HessianResult,
+    HessianUnit,
+    ParameterGradientRequest,
+    ParameterGradientResult,
+    ReferenceEnergyRequest,
+    ReferenceFrequencyRequest,
+    ReferenceHessianRequest,
+    readonly_array,
+)
 from q2mm.models.forcefield import AngleParam, BondParam, ForceField, VdwParam
+from q2mm.models.molecule import Angle, Bond, Molecule, Torsion
 
 if TYPE_CHECKING:
     from q2mm.models.parameters import ParameterLayout
@@ -41,8 +66,8 @@
 _HAS_JAXOPT: bool = importlib.util.find_spec("jaxopt") is not None
 
 # These are populated lazily by ensure_jax() / ensure_jaxopt().
-jax: ModuleType | None = None
-jnp: ModuleType | None = None
+jax: Any = None
+jnp: Any = None
 jaxopt: ModuleType | None = None
 _jax_initialized: bool = False
 
@@ -169,7 +194,7 @@ def match_bond(
 
     """
     matched = forcefield.match_bond(
-        elements,
+        (elements[0], elements[1]),
         env_id=env_id,
         ff_row=ff_row,
         bond_order=bond_order,
@@ -199,7 +224,7 @@ def match_angle(
             ``(None, None)`` if unmatched.
 
     """
-    matched = forcefield.match_angle(elements, env_id=env_id, ff_row=ff_row)
+    matched = forcefield.match_angle((elements[0], elements[1], elements[2]), env_id=env_id, ff_row=ff_row)
     if matched is not None:
         return forcefield.angles.index(matched), matched
     return None, None
@@ -251,6 +276,196 @@ def params_and_coords(
     ensure_jax()
     from q2mm.models.parameters import ParameterLayout
 
-    params = jnp.array(ParameterLayout.from_force_field(forcefield).vector(forcefield), dtype=jnp.float64)  # type: ignore[union-attr]
-    coords = jnp.array(molecule_geometry, dtype=jnp.float64)  # type: ignore[union-attr]
+    params = jnp.array(ParameterLayout.from_force_field(forcefield).vector(forcefield), dtype=jnp.float64)
+    coords = jnp.array(molecule_geometry, dtype=jnp.float64)
     return params, coords
+
+
+@dataclass(frozen=True)
+class MatchedTopology:
+    """Matched topology indices and ForceField parameter maps for JAX backends."""
+
+    matched_bonds: tuple[tuple[Bond, int, BondParam], ...]
+    matched_angles: tuple[tuple[Angle, int, AngleParam], ...]
+    bond_indices: np.ndarray
+    bond_param_map: np.ndarray
+    angle_indices: np.ndarray
+    angle_param_map: np.ndarray
+    torsion_indices: np.ndarray
+    torsion_param_map: np.ndarray
+    atom_vdw_map: np.ndarray
+    explicitly_excluded_atoms: frozenset[int]
+
+
+def match_mm_topology(
+    molecule: Molecule,
+    forcefield: ForceField,
+    *,
+    include_impropers: bool,
+    mask_excluded_vdw: bool,
+) -> MatchedTopology:
+    """Match bonds, angles, torsions, and vdW parameters for a JAX MM state."""
+    matched_bonds: list[tuple[Bond, int, BondParam]] = []
+    for bond in molecule.bonds or ():
+        bond_idx, bond_param = match_bond(
+            forcefield,
+            bond.elements,
+            env_id=bond.env_id,
+            ff_row=bond.ff_row,
+            bond_order=getattr(bond, "bond_order", ""),
+            bond_length=bond.length,
+        )
+        if bond_idx is not None and bond_param is not None:
+            matched_bonds.append((bond, bond_idx, bond_param))
+
+    matched_angles: list[tuple[Angle, int, AngleParam]] = []
+    for angle in molecule.angles or ():
+        angle_idx, angle_param = match_angle(forcefield, angle.elements, env_id=angle.env_id, ff_row=angle.ff_row)
+        if angle_idx is not None and angle_param is not None:
+            matched_angles.append((angle, angle_idx, angle_param))
+
+    torsion_atom_indices: list[tuple[int, int, int, int]] = []
+    torsion_param_map_list: list[int] = []
+    torsion_param_index = {id(p): i for i, p in enumerate(forcefield.torsions)}
+    sources: list[tuple[Sequence[Torsion], bool]] = [(molecule.torsions or (), False)]
+    if include_impropers:
+        sources.append((molecule.improper_torsions, True))
+    for torsion_list, is_improper in sources:
+        for torsion in torsion_list:
+            matches = forcefield.match_torsion(
+                torsion.element_quad, env_id=torsion.env_id, ff_row=torsion.ff_row, is_improper=is_improper
+            )
+            for torsion_param in matches:
+                torsion_atom_indices.append((torsion.atom_i, torsion.atom_j, torsion.atom_k, torsion.atom_l))
+                torsion_param_map_list.append(torsion_param_index[id(torsion_param)])
+
+    atom_types = molecule.atom_types if molecule.atom_types is not None else molecule.symbols
+    excluded_types = {value.casefold() for value in forcefield.nonbonded_excluded_atom_types}
+    explicitly_excluded = frozenset(
+        index
+        for index, (symbol, atom_type) in enumerate(zip(molecule.symbols, atom_types, strict=True))
+        if symbol.casefold() in excluded_types or atom_type.casefold() in excluded_types
+    )
+    atom_vdw_map_list: list[int] = []
+    for index, (symbol, atom_type) in enumerate(zip(molecule.symbols, atom_types, strict=False)):
+        if mask_excluded_vdw and index in explicitly_excluded:
+            atom_vdw_map_list.append(-1)
+            continue
+        vdw_idx, vdw_param = match_vdw(forcefield, atom_type=atom_type, element=symbol)
+        atom_vdw_map_list.append(vdw_idx if vdw_idx is not None and vdw_param is not None else -1)
+
+    unmatched = [i for i, idx in enumerate(atom_vdw_map_list) if idx == -1 and i not in explicitly_excluded]
+    if getattr(forcefield, "vdws", None) and unmatched:
+        raise ValueError(
+            f"Unmatched vdW parameters for atoms at indices {unmatched}. "
+            "Ensure all atom types/elements have corresponding vdW "
+            "parameters in the force field, or remove vdW terms from "
+            "the force field if vdW interactions are not intended."
+        )
+
+    bond_indices = (
+        np.array([(b.atom_i, b.atom_j) for b, _, _ in matched_bonds], dtype=np.int32)
+        if matched_bonds
+        else np.empty((0, 2), dtype=np.int32)
+    )
+    angle_indices = (
+        np.array([(a.atom_i, a.atom_j, a.atom_k) for a, _, _ in matched_angles], dtype=np.int32)
+        if matched_angles
+        else np.empty((0, 3), dtype=np.int32)
+    )
+    torsion_indices = (
+        np.array(torsion_atom_indices, dtype=np.int32) if torsion_atom_indices else np.empty((0, 4), dtype=np.int32)
+    )
+
+    return MatchedTopology(
+        matched_bonds=tuple(matched_bonds),
+        matched_angles=tuple(matched_angles),
+        bond_indices=bond_indices,
+        bond_param_map=np.array([idx for _, idx, _ in matched_bonds], dtype=np.int32),
+        angle_indices=angle_indices,
+        angle_param_map=np.array([idx for _, idx, _ in matched_angles], dtype=np.int32),
+        torsion_indices=torsion_indices,
+        torsion_param_map=np.array(torsion_param_map_list, dtype=np.int32),
+        atom_vdw_map=np.array(atom_vdw_map_list, dtype=np.int32),
+        explicitly_excluded_atoms=explicitly_excluded,
+    )
+
+
+class _PreparedJaxBase(AbstractPreparedBackend):
+    """Shared prepared-session base for JAX and JAX-MD backends."""
+
+    _backend_label: str
+    _backend: Any
+    _state: Any
+
+    def __init__(
+        self,
+        *,
+        backend: Any,
+        info: BackendInfo,
+        case_id: str,
+        molecule: Molecule,
+        force_field: ForceField,
+        layout: ParameterLayout,
+        state: Any,
+    ) -> None:
+        super().__init__(
+            info=info,
+            case_id=case_id,
+            molecule=molecule,
+            force_field=force_field,
+            layout=layout,
+        )
+        self._backend = backend
+        self._state = state
+
+    def _energy(self, request: EnergyRequest | ReferenceEnergyRequest) -> EnergyResult:
+        assert isinstance(request, EnergyRequest)
+        ff = self._ff_for(request.parameters)
+        try:
+            value = self._backend._evaluate_energy(self._state, ff)
+        except Exception as exc:  # noqa: BLE001
+            raise EvaluationError(f"{self._backend_label} energy evaluation failed: {exc}") from exc
+        return EnergyResult(energy=float(value), unit=EnergyUnit.KCAL_PER_MOL, provenance=self._provenance)
+
+    def _raw_hessian(self, ff: ForceField) -> np.ndarray:
+        hess: np.ndarray = self._backend._evaluate_hessian(self._state, ff)
+        return hess
+
+    def _hessian(self, request: HessianRequest | ReferenceHessianRequest) -> HessianResult:
+        assert isinstance(request, HessianRequest)
+        ff = self._ff_for(request.parameters)
+        try:
+            hess = self._raw_hessian(ff)
+        except Exception as exc:  # noqa: BLE001
+            raise EvaluationError(f"{self._backend_label} Hessian evaluation failed: {exc}") from exc
+        return HessianResult(
+            hessian=readonly_array(hess), unit=HessianUnit.HARTREE_PER_BOHR2, provenance=self._provenance
+        )
+
+    def _frequencies(self, request: FrequencyRequest | ReferenceFrequencyRequest) -> FrequencyResult:
+        assert isinstance(request, FrequencyRequest)
+        return self._mm_frequencies(request, backend_label=self._backend_label)
+
+    def _parameter_gradient(self, request: ParameterGradientRequest) -> ParameterGradientResult:
+        ff = self._ff_for(request.parameters)
+        try:
+            energy, grad = self._backend._evaluate_param_grad(self._state, ff)
+        except Exception as exc:  # noqa: BLE001
+            raise EvaluationError(f"{self._backend_label} parameter-gradient evaluation failed: {exc}") from exc
+        return ParameterGradientResult(
+            energy=float(energy),
+            gradient=readonly_array(grad),
+            unit=EnergyUnit.KCAL_PER_MOL,
+            provenance=self._provenance,
+        )
+
+    def _batched_energy(self, request: BatchedEnergyRequest) -> BatchedEnergyResult:
+        mat = self._validate_matrix(request.parameter_matrix)
+        try:
+            energies = self._backend._evaluate_batched_energy(self._state, self.force_field, mat)
+        except Exception as exc:  # noqa: BLE001
+            raise EvaluationError(f"{self._backend_label} batched-energy evaluation failed: {exc}") from exc
+        return BatchedEnergyResult(
+            energies=readonly_array(energies), unit=EnergyUnit.KCAL_PER_MOL, provenance=self._provenance
+        )
diff --git a/q2mm/backends/mm/_term_support.py b/q2mm/backends/mm/_term_support.py
index e02c2841..8570e137 100644
--- a/q2mm/backends/mm/_term_support.py
+++ b/q2mm/backends/mm/_term_support.py
@@ -5,7 +5,7 @@
 import re
 from typing import Literal
 
-from q2mm.backends.contracts import PreparationError
+from q2mm.backends.contracts import BackendInfo, PreparationError, PreparationRequest
 from q2mm.models.forcefield import ForceField, TorsionParam
 
 _Term = Literal[
@@ -47,3 +47,23 @@ def _validate_term_support(force_field: ForceField, *, backend: str, unsupported
             f"{backend} does not support populated {', '.join(rejected)} "
             f"for functional form {force_field.functional_form.value!r}."
         )
+
+
+def _require_mm_force_field(
+    request: PreparationRequest,
+    info: BackendInfo,
+    *,
+    backend: str,
+    unsupported: frozenset[_Term] = frozenset(),
+) -> ForceField:
+    """Require a base ForceField with a supported functional form and supported populated terms."""
+    if request.force_field is None:
+        raise PreparationError(f"{backend} requires a base ForceField in the PreparationRequest.")
+    form = request.force_field.functional_form.value
+    if not info.supports_form(form):
+        raise PreparationError(
+            f"{backend} does not support functional form {form!r}. Supported: {sorted(info.functional_forms)}"
+        )
+    if unsupported:
+        _validate_term_support(request.force_field, backend=backend, unsupported=unsupported)
+    return request.force_field
diff --git a/q2mm/backends/mm/batched.py b/q2mm/backends/mm/batched.py
index c06cf3b6..bb0b4332 100644
--- a/q2mm/backends/mm/batched.py
+++ b/q2mm/backends/mm/batched.py
@@ -12,16 +12,22 @@
 from __future__ import annotations
 
 import logging
-from typing import TYPE_CHECKING
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 logger = logging.getLogger(__name__)
 
+jax: Any = None
+jnp: Any = None
+
 try:
-    import jax
-    import jax.numpy as jnp
+    import jax as _jax_mod
+    import jax.numpy as _jnp_mod
 
+    jax = _jax_mod
+    jnp = _jnp_mod
     _HAS_JAX = True
 except ImportError:  # pragma: no cover
     _HAS_JAX = False
@@ -31,6 +37,7 @@
     BatchedHessianResult,
     EvaluationError,
     HessianUnit,
+    PreparedBackend,
     readonly_array,
 )
 from q2mm.models.units import KCALMOLA2_TO_HESSIAN_AU
@@ -173,22 +180,28 @@ def _batched_hessians(self, params_np: np.ndarray) -> list[np.ndarray]:
         params = jnp.array(params_np, dtype=jnp.float64)
 
         if state._coord_hess_fn is None:
+            energy_fn = state._energy_fn
+            assert energy_fn is not None
 
-            def _energy_of_flat_coords(flat_coords: jnp.ndarray, params_: jnp.ndarray) -> jnp.ndarray:
-                return state._energy_fn(params_, flat_coords.reshape(-1, 3))
+            def _energy_of_flat_coords(flat_coords: Any, params_: Any) -> Any:
+                return energy_fn(params_, flat_coords.reshape(-1, 3))
 
             state._coord_hess_fn = jax.jit(jax.hessian(_energy_of_flat_coords, argnums=0))
 
+        coord_hess_fn = state._coord_hess_fn
+        assert coord_hess_fn is not None
         if len(self._geometries) == 1:
             flat = jnp.array(self._geometries[0], dtype=jnp.float64).flatten()
-            hess = state._coord_hess_fn(flat, params)
+            hess = coord_hess_fn(flat, params)
             return [np.asarray(hess) * KCALMOLA2_TO_HESSIAN_AU]
 
         if state._batched_coord_hess_fn is None:
-            state._batched_coord_hess_fn = jax.jit(jax.vmap(state._coord_hess_fn, in_axes=(0, None)))
+            state._batched_coord_hess_fn = jax.jit(jax.vmap(coord_hess_fn, in_axes=(0, None)))
 
+        batched_coord_hess_fn = state._batched_coord_hess_fn
+        assert batched_coord_hess_fn is not None
         batch_coords = jnp.stack([jnp.array(g, dtype=jnp.float64).flatten() for g in self._geometries])
-        batch_hess = state._batched_coord_hess_fn(batch_coords, params)
+        batch_hess = batched_coord_hess_fn(batch_coords, params)
         return [np.asarray(batch_hess[i]) * KCALMOLA2_TO_HESSIAN_AU for i in range(len(self._geometries))]
 
 
@@ -197,7 +210,7 @@ def _energy_of_flat_coords(flat_coords: jnp.ndarray, params_: jnp.ndarray) -> jn
 # ---------------------------------------------------------------------------
 
 
-def group_by_topology(sessions: list[PreparedJax]) -> list[PreparedJaxBatch]:
+def group_by_topology(sessions: Sequence[PreparedBackend]) -> list[PreparedJaxBatch]:
     """Group prepared sessions into typed batches by topology compatibility.
 
     Two sessions are compatible if all their static kernel inputs and
@@ -214,8 +227,12 @@ def group_by_topology(sessions: list[PreparedJax]) -> list[PreparedJaxBatch]:
         List of :class:`PreparedJaxBatch`, one per distinct topology.
 
     """
+    from q2mm.backends.mm.jax_engine import PreparedJax
+
     groups: dict[str, list[PreparedJax]] = {}
     for session in sessions:
+        if not isinstance(session, PreparedJax):
+            raise EvaluationError(f"PreparedJaxBatch requires PreparedJax sessions; got {type(session).__name__}.")
         sig = _topology_signature(session._state)
         groups.setdefault(sig, []).append(session)
     return [PreparedJaxBatch(group) for group in groups.values()]
diff --git a/q2mm/backends/mm/jax_engine.py b/q2mm/backends/mm/jax_engine.py
index 214ae7ee..a37377bf 100644
--- a/q2mm/backends/mm/jax_engine.py
+++ b/q2mm/backends/mm/jax_engine.py
@@ -26,7 +26,7 @@
 import math
 from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -35,37 +35,25 @@
 logger = logging.getLogger(__name__)
 
 from q2mm.backends.contracts import (
-    AbstractPreparedBackend,
     BackendInfo,
     BackendProvenance,
     BackendRole,
     BackendUnavailableError,
-    BatchedEnergyRequest,
-    BatchedEnergyResult,
     Capability,
-    EnergyRequest,
-    EnergyResult,
     EnergyUnit,
     EvaluationError,
-    FrequencyRequest,
-    FrequencyResult,
-    FrequencyUnit,
     GeometryResult,
     HessianJacobianRequest,
     HessianJacobianResult,
-    HessianRequest,
-    HessianResult,
     HessianUnit,
     LengthUnit,
     MinimizationRequest,
-    ParameterGradientRequest,
-    ParameterGradientResult,
     PreparationError,
     PreparationRequest,
     PreparedBackend,
     readonly_array,
 )
-from q2mm.backends.mm._term_support import _validate_term_support
+from q2mm.backends.mm._term_support import _require_mm_force_field
 from q2mm.models.units import KCALMOLA2_TO_HESSIAN_AU
 from q2mm.constants import (
     MM3_BOND_C3,
@@ -83,20 +71,22 @@
 
 from q2mm.backends.mm._jax_common import (
     _HAS_JAX,
+    _PreparedJaxBase,
     ensure_jax as _ensure_jax_common,
     layout_block_offsets,
-    match_angle as _match_angle,
-    match_bond as _match_bond,
-    match_vdw as _match_vdw,
+    match_mm_topology,
     params_and_coords as _params_and_coords_impl,
 )
 
 if TYPE_CHECKING:
     from q2mm.backends.mm.batched import PreparedJaxBatch
+    from q2mm.models.parameters import ParameterLayout
+
+type JaxArray = Any
 
 # Lazy placeholders — populated by _ensure_jax().
-jax = None
-jnp = None
+jax: Any = None
+jnp: Any = None
 
 
 def _ensure_jax() -> None:
@@ -128,7 +118,7 @@ def _ensure_jax() -> None:
 # ---------------------------------------------------------------------------
 
 
-def _safe_norm(x: jnp.ndarray, axis: int = -1) -> jnp.ndarray:
+def _safe_norm(x: JaxArray, axis: int = -1) -> JaxArray:
     """Euclidean norm with gradient-safe floor to avoid NaN at zero.
 
     Args:
@@ -142,7 +132,7 @@ def _safe_norm(x: jnp.ndarray, axis: int = -1) -> jnp.ndarray:
     return jnp.sqrt(jnp.maximum(jnp.sum(x**2, axis=axis), 1e-30))
 
 
-def _angle_from_vectors(a: jnp.ndarray, b: jnp.ndarray) -> jnp.ndarray:
+def _angle_from_vectors(a: JaxArray, b: JaxArray) -> JaxArray:
     """Angle between two vectors via ``atan2(|a×b|, a·b)`` with safe VJP.
 
     Returns the angle (radians) in ``[0, π]`` with **correct gradients
@@ -184,7 +174,7 @@ def _angle_from_vectors(a: jnp.ndarray, b: jnp.ndarray) -> jnp.ndarray:
 
 
 @functools.lru_cache(maxsize=1)
-def _angle_from_vectors_impl() -> Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]:
+def _angle_from_vectors_impl() -> Callable[[JaxArray, JaxArray], JaxArray]:
     """Build the jax.custom_vjp-decorated angle function lazily.
 
     The decorator can't be applied at module import because ``jax`` is
@@ -195,7 +185,7 @@ def _angle_from_vectors_impl() -> Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarr
     _ensure_jax()
 
     @jax.custom_vjp
-    def _angle(a: jnp.ndarray, b: jnp.ndarray) -> jnp.ndarray:
+    def _angle(a: JaxArray, b: JaxArray) -> JaxArray:
         na = jnp.sqrt(jnp.sum(a * a, axis=-1) + 1e-30)
         nb = jnp.sqrt(jnp.sum(b * b, axis=-1) + 1e-30)
         cross = jnp.cross(a, b, axis=-1)
@@ -228,7 +218,8 @@ def _bwd(res, g):  # noqa: ANN001, ANN202
         return (g[..., None] * da, g[..., None] * db)
 
     _angle.defvjp(_fwd, _bwd)
-    return _angle
+    fn: Callable[[JaxArray, JaxArray], JaxArray] = _angle
+    return fn
 
 
 # ---------------------------------------------------------------------------
@@ -236,9 +227,7 @@ def _bwd(res, g):  # noqa: ANN001, ANN202
 # ---------------------------------------------------------------------------
 
 
-def _harmonic_bond_energy(
-    k: jnp.ndarray, r0: jnp.ndarray, coords: jnp.ndarray, bond_indices: jnp.ndarray
-) -> jnp.ndarray:
+def _harmonic_bond_energy(k: JaxArray, r0: JaxArray, coords: JaxArray, bond_indices: JaxArray) -> JaxArray:
     """Harmonic bond stretch: ``E = Σ k_i · (r_i − r0_i)²``.
 
     Args:
@@ -260,9 +249,7 @@ def _harmonic_bond_energy(
     return jnp.sum(k * (r - r0) ** 2)
 
 
-def _harmonic_angle_energy(
-    k: jnp.ndarray, theta0: jnp.ndarray, coords: jnp.ndarray, angle_indices: jnp.ndarray
-) -> jnp.ndarray:
+def _harmonic_angle_energy(k: JaxArray, theta0: JaxArray, coords: JaxArray, angle_indices: JaxArray) -> JaxArray:
     """Harmonic angle bend: ``E = Σ k_i · (θ_i − θ0_i)²``.
 
     Args:
@@ -286,8 +273,8 @@ def _harmonic_angle_energy(
 
 
 def _lj_12_6_energy(
-    per_atom_sigma: jnp.ndarray, per_atom_epsilon: jnp.ndarray, coords: jnp.ndarray, pair_indices: jnp.ndarray
-) -> jnp.ndarray:
+    per_atom_sigma: JaxArray, per_atom_epsilon: JaxArray, coords: JaxArray, pair_indices: JaxArray
+) -> JaxArray:
     """Lennard-Jones 12-6 energy with geometric combining rules.
 
     ``E = Σ_{i<j} 4·ε_ij · [(σ_ij/r)¹² − (σ_ij/r)⁶]`` where
@@ -333,7 +320,7 @@ def _lj_12_6_energy(
 _RAD_TO_DEG = 180.0 / 3.141592653589793
 
 
-def _mm3_bond_energy(k: jnp.ndarray, r0: jnp.ndarray, coords: jnp.ndarray, bond_indices: jnp.ndarray) -> jnp.ndarray:
+def _mm3_bond_energy(k: JaxArray, r0: JaxArray, coords: JaxArray, bond_indices: JaxArray) -> JaxArray:
     """MM3 cubic bond stretch: ``E = Σ k·dr²·(1 − 2.55·dr + c4·dr²)``.
 
     Args:
@@ -352,9 +339,7 @@ def _mm3_bond_energy(k: jnp.ndarray, r0: jnp.ndarray, coords: jnp.ndarray, bond_
     return jnp.sum(k * dr**2 * (1.0 - MM3_BOND_C3 * dr + MM3_BOND_C4 * dr**2))
 
 
-def _mm3_angle_energy(
-    k: jnp.ndarray, theta0: jnp.ndarray, coords: jnp.ndarray, angle_indices: jnp.ndarray
-) -> jnp.ndarray:
+def _mm3_angle_energy(k: JaxArray, theta0: JaxArray, coords: JaxArray, angle_indices: JaxArray) -> JaxArray:
     """MM3 sextic angle bend with degree-based anharmonic corrections.
 
     ``E = Σ k·dθ²·(1 + a3·dθ_deg + a4·dθ_deg² + a5·dθ_deg³ + a6·dθ_deg⁴)``
@@ -389,11 +374,11 @@ def _mm3_angle_energy(
 
 
 def _mm3_vdw_energy(
-    per_atom_radius: jnp.ndarray,
-    per_atom_epsilon: jnp.ndarray,
-    coords: jnp.ndarray,
-    pair_indices: jnp.ndarray,
-) -> jnp.ndarray:
+    per_atom_radius: JaxArray,
+    per_atom_epsilon: JaxArray,
+    coords: JaxArray,
+    pair_indices: JaxArray,
+) -> JaxArray:
     """MM3 Buckingham exp-6 vdW with short-range repulsive wall.
 
     ``E = ε·(184000·exp(−12·r/rv) − 2.25·(rv/r)⁶)``
@@ -449,11 +434,11 @@ def _mm3_vdw_energy(
 
 
 def _mm3_dipole_energy(
-    dipole_moments: jnp.ndarray,
-    coords: jnp.ndarray,
-    bond_indices: jnp.ndarray,
-    dipole_pair_indices: jnp.ndarray,
-) -> jnp.ndarray:
+    dipole_moments: JaxArray,
+    coords: JaxArray,
+    bond_indices: JaxArray,
+    dipole_pair_indices: JaxArray,
+) -> JaxArray:
     """MM3 bond-dipole electrostatic energy.
 
     Each bond with a non-zero dipole moment contributes a point dipole
@@ -519,13 +504,13 @@ def _mm3_dipole_energy(
 
 
 def _mm3_stretch_bend_energy(
-    k_sb: jnp.ndarray,
-    r0_ij: jnp.ndarray,
-    r0_jk: jnp.ndarray,
-    theta0: jnp.ndarray,
-    coords: jnp.ndarray,
-    angle_indices: jnp.ndarray,
-) -> jnp.ndarray:
+    k_sb: JaxArray,
+    r0_ij: JaxArray,
+    r0_jk: JaxArray,
+    theta0: JaxArray,
+    coords: JaxArray,
+    angle_indices: JaxArray,
+) -> JaxArray:
     """MM3 stretch-bend cross-term energy.
 
     ``E = Σ k_sb · [(r_ij − r0_ij) + (r_jk − r0_jk)] · (θ − θ0)``
@@ -568,7 +553,7 @@ def _mm3_stretch_bend_energy(
     return jnp.sum(w * k_sb * dr_sum * dtheta)
 
 
-def _smoothstep(x: jnp.ndarray, edge0: float, edge1: float) -> jnp.ndarray:
+def _smoothstep(x: JaxArray, edge0: float, edge1: float) -> JaxArray:
     """Hermite smoothstep: 0 below *edge0*, 1 above *edge1*, C¹ between."""
     t = jnp.clip((x - edge0) / (edge1 - edge0), 0.0, 1.0)
     return t * t * (3.0 - 2.0 * t)
@@ -588,12 +573,12 @@ def _smoothstep(x: jnp.ndarray, edge0: float, edge1: float) -> jnp.ndarray:
 
 
 def _torsion_energy(
-    k: jnp.ndarray,
-    periodicity: jnp.ndarray,
-    phase: jnp.ndarray,
-    coords: jnp.ndarray,
-    torsion_indices: jnp.ndarray,
-) -> jnp.ndarray:
+    k: JaxArray,
+    periodicity: JaxArray,
+    phase: JaxArray,
+    coords: JaxArray,
+    torsion_indices: JaxArray,
+) -> JaxArray:
     """Cosine torsion energy: ``E = Σ w_i · k_i · (1 + cos(n_i·φ_i − γ_i))``.
 
     Uses an atan2-based signed dihedral angle computation.  Torsion terms
@@ -766,6 +751,7 @@ class _JaxState:
     _kernel_signature: str | None = field(default=None, repr=False)
     # Compiled energy function (captures topology, JIT-compiled)
     _energy_fn: Callable | None = field(default=None, repr=False)
+    _batched_energy_fn: Callable | None = field(default=None, repr=False)
     _grad_fn: Callable | None = field(default=None, repr=False)
     _coord_hess_fn: Callable | None = field(default=None, repr=False)
     _batched_coord_hess_fn: Callable | None = field(default=None, repr=False)
@@ -851,16 +837,13 @@ def prepare(self, request: PreparationRequest) -> PreparedJax:
         """
         from q2mm.models.parameters import ParameterLayout
 
+        info = self.info
         if request.force_field is None:
             raise PreparationError("JAX requires a base ForceField in the PreparationRequest.")
-        info = self.info
         form = request.force_field.functional_form.value
-        if not info.supports_form(form):
-            raise PreparationError(
-                f"JAX does not support functional form {form!r}. Supported: {sorted(info.functional_forms)}"
-            )
-        _validate_term_support(
-            request.force_field,
+        force_field = _require_mm_force_field(
+            request,
+            info,
             backend="JAX",
             unsupported=frozenset(
                 {"CMAP", "vdW reduction", "wildcard torsions", "stretch-bend", "bond dipoles"}
@@ -868,9 +851,9 @@ def prepare(self, request: PreparationRequest) -> PreparedJax:
                 else {"CMAP", "vdW reduction", "wildcard torsions"}
             ),
         )
-        layout = ParameterLayout.from_force_field(request.force_field)
+        layout = ParameterLayout.from_force_field(force_field)
         try:
-            state = self._build_state(request.molecule, request.force_field)
+            state = self._build_state(request.molecule, force_field)
         except Exception as exc:  # noqa: BLE001
             raise PreparationError(f"JAX failed to prepare case {request.case_id!r}: {exc}") from exc
         return PreparedJax(
@@ -878,7 +861,7 @@ def prepare(self, request: PreparationRequest) -> PreparedJax:
             info=info,
             case_id=request.case_id,
             molecule=request.molecule,
-            force_field=request.force_field,
+            force_field=force_field,
             layout=layout,
             state=state,
         )
@@ -892,7 +875,7 @@ def prepare_hessian_batches(self, sessions: Sequence[PreparedBackend]) -> list[P
         """
         from q2mm.backends.mm.batched import group_by_topology
 
-        return group_by_topology(list(sessions))
+        return group_by_topology(sessions)
 
     def _build_state(self, molecule: Molecule, forcefield: ForceField) -> _JaxState:
         """Build a compiled private native state for *molecule* + *forcefield*.
@@ -900,48 +883,21 @@ def _build_state(self, molecule: Molecule, forcefield: ForceField) -> _JaxState:
         Requires an explicit :class:`ForceField`; no auto-generation and no
         state pass-through union.
         """
-        # Match bonds
-        bond_atom_indices = []
-        bond_param_map = []
-        for bond in molecule.bonds:
-            idx, param = _match_bond(
-                forcefield,
-                bond.elements,
-                env_id=bond.env_id,
-                ff_row=bond.ff_row,
-                bond_order=getattr(bond, "bond_order", ""),
-                bond_length=bond.length,
-            )
-            if param is not None:
-                bond_atom_indices.append((bond.atom_i, bond.atom_j))
-                bond_param_map.append(idx)
-
-        # Match angles
-        angle_atom_indices = []
-        angle_param_map = []
-        for angle in molecule.angles:
-            idx, param = _match_angle(forcefield, angle.elements, env_id=angle.env_id, ff_row=angle.ff_row)
-            if param is not None:
-                angle_atom_indices.append((angle.atom_i, angle.atom_j, angle.atom_k))
-                angle_param_map.append(idx)
+        topo = match_mm_topology(molecule, forcefield, include_impropers=True, mask_excluded_vdw=False)
 
         # Match Urey-Bradley terms (1-3 distance for CHARMM angles)
         ub_angles = forcefield._ub_angles
         ub_angle_to_ub_idx = {id(a): i for i, a in enumerate(ub_angles)}
         ub_atom_indices: list[tuple[int, int]] = []
         ub_param_map: list[int] = []
-        for angle in molecule.angles:
-            matched = forcefield.match_angle(angle.elements, env_id=angle.env_id, ff_row=angle.ff_row)
-            if matched is not None and matched.ub_force_constant is not None and matched.ub_equilibrium is not None:
-                ub_idx = ub_angle_to_ub_idx.get(id(matched))
+        for angle, _angle_idx, matched_angle in topo.matched_angles:
+            if matched_angle.ub_force_constant is not None and matched_angle.ub_equilibrium is not None:
+                ub_idx = ub_angle_to_ub_idx.get(id(matched_angle))
                 if ub_idx is not None:
                     ub_atom_indices.append((angle.atom_i, angle.atom_k))
                     ub_param_map.append(ub_idx)
 
-        # Match stretch-bend cross terms
-        # Use match_stretch_bend() with proper env_id/ff_row fallback.
-        # Look up adjacent bonds from the already-matched bond set (not
-        # re-matching with the wrong env_id).
+        # Match stretch-bend cross terms using already-matched bonds and angles.
         sb_angle_indices: list[tuple[int, int, int]] = []
         sb_param_map_list: list[int] = []
         sb_bond_ij_idx: list[int] = []
@@ -949,114 +905,39 @@ def _build_state(self, molecule: Molecule, forcefield: ForceField) -> _JaxState:
         sb_angle_idx: list[int] = []
         sb_param_index = {id(p): i for i, p in enumerate(forcefield.stretch_bends)}
         bond_ff_index = {id(p): i for i, p in enumerate(forcefield.bonds)}
-        angle_ff_index = {id(p): i for i, p in enumerate(forcefield.angles)}
 
-        # Pre-build lookup from atom pair → matched bond param (using the
-        # bond's own env_id, not the angle's).
         matched_bond_by_atoms: dict[tuple[int, int], BondParam] = {}
-        for bond in molecule.bonds:
-            _, param = _match_bond(
-                forcefield,
-                bond.elements,
-                env_id=bond.env_id,
-                ff_row=bond.ff_row,
-                bond_order=getattr(bond, "bond_order", ""),
-                bond_length=bond.length,
-            )
-            if param is not None:
-                matched_bond_by_atoms[(bond.atom_i, bond.atom_j)] = param
-                matched_bond_by_atoms[(bond.atom_j, bond.atom_i)] = param
+        for bond, _bond_idx, bond_param in topo.matched_bonds:
+            matched_bond_by_atoms[(bond.atom_i, bond.atom_j)] = bond_param
+            matched_bond_by_atoms[(bond.atom_j, bond.atom_i)] = bond_param
 
-        for angle in molecule.angles:
+        for angle, angle_idx, _angle_param in topo.matched_angles:
             sb = forcefield.match_stretch_bend(angle.elements, env_id=angle.env_id, ff_row=angle.ff_row)
             if sb is None:
                 continue
             bond_ij = matched_bond_by_atoms.get((angle.atom_i, angle.atom_j))
             bond_jk = matched_bond_by_atoms.get((angle.atom_j, angle.atom_k))
-            angle_param = forcefield.match_angle(angle.elements, env_id=angle.env_id, ff_row=angle.ff_row)
-            if bond_ij is not None and bond_jk is not None and angle_param is not None:
+            if bond_ij is not None and bond_jk is not None:
                 sb_angle_indices.append((angle.atom_i, angle.atom_j, angle.atom_k))
                 sb_param_map_list.append(sb_param_index[id(sb)])
                 sb_bond_ij_idx.append(bond_ff_index[id(bond_ij)])
                 sb_bond_jk_idx.append(bond_ff_index[id(bond_jk)])
-                sb_angle_idx.append(angle_ff_index[id(angle_param)])
-
-        # Match torsions — each detected torsion may match multiple FF
-        # entries (one per periodicity component), each becoming a separate
-        # term.  Both proper and improper torsions are routed through the
-        # same cosine-series energy term (mirrors OpenMM, which puts both
-        # into a single PeriodicTorsionForce; see openmm.py:1080-1132).
-        torsion_atom_indices: list[tuple[int, int, int, int]] = []
-        torsion_param_map: list[int] = []
-        torsion_param_index = {id(p): i for i, p in enumerate(forcefield.torsions)}
-        for torsion in molecule.torsions:
-            matches = forcefield.match_torsion(
-                torsion.element_quad, env_id=torsion.env_id, ff_row=torsion.ff_row, is_improper=False
-            )
-            for param in matches:
-                j_ff = torsion_param_index[id(param)]
-                torsion_atom_indices.append((torsion.atom_i, torsion.atom_j, torsion.atom_k, torsion.atom_l))
-                torsion_param_map.append(j_ff)
-        for imp in molecule.improper_torsions:
-            matches = forcefield.match_torsion(imp.element_quad, env_id=imp.env_id, ff_row=imp.ff_row, is_improper=True)
-            for param in matches:
-                j_ff = torsion_param_index[id(param)]
-                torsion_atom_indices.append((imp.atom_i, imp.atom_j, imp.atom_k, imp.atom_l))
-                torsion_param_map.append(j_ff)
-
-        # Match vdW
-        atom_vdw_map = []
-        for atom_index, (symbol, atom_type) in enumerate(zip(molecule.symbols, molecule.atom_types, strict=False)):
-            idx, param = _match_vdw(forcefield, atom_type=atom_type, element=symbol)
-            if param is not None:
-                atom_vdw_map.append(idx)
-            else:
-                atom_vdw_map.append(-1)
-
-        # Validate vdW mapping: if the force field defines vdW parameters,
-        # disallow any unmatched atoms.  Using -1 as an index would silently
-        # select the last vdW entry via JAX negative indexing, corrupting
-        # the energy and gradients.  When the force field defines no vdW
-        # terms at all, vdW energy is effectively disabled and unmatched
-        # atoms are safe.
-        excluded_types = {value.casefold() for value in forcefield.nonbonded_excluded_atom_types}
-        explicitly_excluded = frozenset(
-            index
-            for index, (symbol, atom_type) in enumerate(zip(molecule.symbols, molecule.atom_types, strict=True))
-            if symbol.casefold() in excluded_types or atom_type.casefold() in excluded_types
-        )
-        unmatched = [i for i, idx in enumerate(atom_vdw_map) if idx == -1 and i not in explicitly_excluded]
-        if getattr(forcefield, "vdws", None) and unmatched:
-            raise ValueError(
-                f"Unmatched vdW parameters for atoms at indices {unmatched}. "
-                "Ensure all atom types/elements have corresponding vdW "
-                "parameters in the force field, or remove vdW terms from "
-                "the force field if vdW interactions are not intended."
-            )
+                sb_angle_idx.append(angle_idx)
 
         # Build vdW pair list with 1-2 and 1-3 exclusions
         vdw_pairs = _build_vdw_pairs(
             len(molecule.symbols),
-            [(b.atom_i, b.atom_j) for b in molecule.bonds],
-            excluded_atoms=explicitly_excluded,
+            [(b.atom_i, b.atom_j) for b in molecule.bonds or ()],
+            excluded_atoms=topo.explicitly_excluded_atoms,
         )
 
-        # Build bond-dipole electrostatics topology
-        # Collect bonds with non-zero dipole moments and their atom indices
+        # Build bond-dipole electrostatics topology from matched bonds
         dipole_bond_list: list[tuple[int, int]] = []
         dipole_moment_list: list[float] = []
-        for bond_idx, bond in enumerate(molecule.bonds):
-            idx, param = _match_bond(
-                forcefield,
-                bond.elements,
-                env_id=bond.env_id,
-                ff_row=bond.ff_row,
-                bond_order=getattr(bond, "bond_order", ""),
-                bond_length=bond.length,
-            )
-            if param is not None and abs(param.dipole_moment) > 1e-10:
+        for bond, _bond_idx, bond_param in topo.matched_bonds:
+            if abs(bond_param.dipole_moment) > 1e-10:
                 dipole_bond_list.append((bond.atom_i, bond.atom_j))
-                dipole_moment_list.append(param.dipole_moment)
+                dipole_moment_list.append(bond_param.dipole_moment)
 
         # Build dipole pair list: exclude bond pairs that share an atom
         dipole_pairs_list: list[tuple[int, int]] = []
@@ -1064,7 +945,6 @@ def _build_state(self, molecule: Molecule, forcefield: ForceField) -> _JaxState:
             bond_atom_sets = [set(pair) for pair in dipole_bond_list]
             for i in range(len(dipole_bond_list)):
                 for j in range(i + 1, len(dipole_bond_list)):
-                    # Exclude if the two bonds share any atom (1-2/1-3 at bond level)
                     if not bond_atom_sets[i] & bond_atom_sets[j]:
                         dipole_pairs_list.append((i, j))
 
@@ -1075,30 +955,20 @@ def _build_state(self, molecule: Molecule, forcefield: ForceField) -> _JaxState:
         dipole_pair_arr = (
             np.array(dipole_pairs_list, dtype=np.int32) if dipole_pairs_list else np.empty((0, 2), dtype=np.int32)
         )
-
-        bond_indices_arr = (
-            np.array(bond_atom_indices, dtype=np.int32) if bond_atom_indices else np.empty((0, 2), dtype=np.int32)
-        )
-        angle_indices_arr = (
-            np.array(angle_atom_indices, dtype=np.int32) if angle_atom_indices else np.empty((0, 3), dtype=np.int32)
-        )
-        torsion_indices_arr = (
-            np.array(torsion_atom_indices, dtype=np.int32) if torsion_atom_indices else np.empty((0, 4), dtype=np.int32)
-        )
         ub_indices_arr = (
             np.array(ub_atom_indices, dtype=np.int32) if ub_atom_indices else np.empty((0, 2), dtype=np.int32)
         )
 
         state = _JaxState(
             molecule=molecule,
-            bond_indices=bond_indices_arr,
-            angle_indices=angle_indices_arr,
-            torsion_indices=torsion_indices_arr,
+            bond_indices=topo.bond_indices,
+            angle_indices=topo.angle_indices,
+            torsion_indices=topo.torsion_indices,
             vdw_pair_indices=vdw_pairs,
-            bond_param_map=np.array(bond_param_map, dtype=np.int32),
-            angle_param_map=np.array(angle_param_map, dtype=np.int32),
-            torsion_param_map=np.array(torsion_param_map, dtype=np.int32),
-            atom_vdw_map=np.array(atom_vdw_map, dtype=np.int32),
+            bond_param_map=topo.bond_param_map,
+            angle_param_map=topo.angle_param_map,
+            torsion_param_map=topo.torsion_param_map,
+            atom_vdw_map=topo.atom_vdw_map,
             n_bond_types=len(forcefield.bonds),
             n_angle_types=len(forcefield.angles),
             n_torsion_types=len(forcefield.torsions),
@@ -1126,13 +996,14 @@ def _build_state(self, molecule: Molecule, forcefield: ForceField) -> _JaxState:
         state._energy_fn = _compile_energy_fn(state, forcefield)
         return state
 
-    def _params_and_coords(self, state: _JaxState, forcefield: ForceField) -> tuple[jnp.ndarray, jnp.ndarray]:
+    def _params_and_coords(self, state: _JaxState, forcefield: ForceField) -> tuple[JaxArray, JaxArray]:
         """Extract JAX arrays from force field and molecule."""
         return _params_and_coords_impl(state.molecule.geometry, forcefield)
 
     def _evaluate_energy(self, state: _JaxState, forcefield: ForceField) -> float:
         """Potential energy in kcal/mol for a prepared native state."""
         params, coords = self._params_and_coords(state, forcefield)
+        assert state._energy_fn is not None
         return float(state._energy_fn(params, coords))
 
     def _evaluate_param_grad(self, state: _JaxState, forcefield: ForceField) -> tuple[float, np.ndarray]:
@@ -1152,7 +1023,7 @@ def _evaluate_batched_energy(
         """Evaluate a batch of parameter vectors via ``jax.vmap`` (kcal/mol)."""
         coords = jnp.array(state.molecule.geometry, dtype=jnp.float64)
         batch_params = jnp.array(param_matrix, dtype=jnp.float64)
-        if not hasattr(state, "_batched_energy_fn") or state._batched_energy_fn is None:
+        if state._batched_energy_fn is None:
             state._batched_energy_fn = jax.jit(jax.vmap(state._energy_fn, in_axes=(0, None)))
         return np.asarray(state._batched_energy_fn(batch_params, coords))
 
@@ -1160,9 +1031,11 @@ def _evaluate_hessian(self, state: _JaxState, forcefield: ForceField) -> np.ndar
         """Coordinate Hessian d2E/dcoords2 in Hartree/Bohr2."""
         params, coords = self._params_and_coords(state, forcefield)
         if state._coord_hess_fn is None:
+            energy_fn = state._energy_fn
+            assert energy_fn is not None
 
-            def _energy_of_flat_coords(flat_coords: jnp.ndarray, params_: jnp.ndarray) -> jnp.ndarray:
-                return state._energy_fn(params_, flat_coords.reshape(-1, 3))
+            def _energy_of_flat_coords(flat_coords: JaxArray, params_: JaxArray) -> JaxArray:
+                return energy_fn(params_, flat_coords.reshape(-1, 3))
 
             state._coord_hess_fn = jax.jit(jax.hessian(_energy_of_flat_coords, argnums=0))
 
@@ -1176,13 +1049,15 @@ def _evaluate_hessian_param_jac(self, state: _JaxState, forcefield: ForceField)
         flat_coords = coords.flatten()
 
         if state._hess_param_jac_fn is None:
+            energy_fn = state._energy_fn
+            assert energy_fn is not None
 
-            def _energy_of_flat_coords(flat_coords_: jnp.ndarray, params_: jnp.ndarray) -> jnp.ndarray:
-                return state._energy_fn(params_, flat_coords_.reshape(-1, 3))
+            def _energy_of_flat_coords(flat_coords_: JaxArray, params_: JaxArray) -> JaxArray:
+                return energy_fn(params_, flat_coords_.reshape(-1, 3))
 
             hess_fn = jax.hessian(_energy_of_flat_coords, argnums=0)
 
-            def _hess_and_jac(flat_coords_: jnp.ndarray, params_: jnp.ndarray) -> tuple:
+            def _hess_and_jac(flat_coords_: JaxArray, params_: JaxArray) -> tuple:
                 h = hess_fn(flat_coords_, params_)
                 dh_dp = jax.jacrev(hess_fn, argnums=1)(flat_coords_, params_)
                 return h, dh_dp
@@ -1208,6 +1083,7 @@ def _evaluate_minimize(
 
         params, coords = self._params_and_coords(state, forcefield)
         energy_fn = state._energy_fn
+        assert energy_fn is not None
         coord_grad_fn = jax.jit(jax.grad(lambda c, p: energy_fn(p, c.reshape(-1, 3)), argnums=0))
 
         x0 = np.asarray(coords.flatten())
@@ -1236,7 +1112,7 @@ def gradient(x: np.ndarray) -> np.ndarray:
         return opt_energy, list(state.molecule.symbols), opt_coords
 
 
-class PreparedJax(AbstractPreparedBackend):
+class PreparedJax(_PreparedJaxBase):
     """Prepared JAX session for a single training case.
 
     Owns the molecule, base force field, parameter layout, and one compiled
@@ -1246,6 +1122,10 @@ class PreparedJax(AbstractPreparedBackend):
     cases.  Compatible conformers may share only the compiled kernel.
     """
 
+    _backend_label = "JAX"
+    _backend: JaxBackend
+    _state: _JaxState
+
     def __init__(
         self,
         *,
@@ -1254,22 +1134,18 @@ def __init__(
         case_id: str,
         molecule: Molecule,
         force_field: ForceField,
-        layout: object,
+        layout: ParameterLayout,
         state: _JaxState,
     ) -> None:
         super().__init__(
+            backend=backend,
             info=info,
             case_id=case_id,
             molecule=molecule,
             force_field=force_field,
             layout=layout,
+            state=state,
         )
-        self._backend = backend
-        self._state = state
-
-    def _ff_for(self, parameters: np.ndarray) -> ForceField:
-        vec = self._validate_vector(parameters)
-        return self.layout.replace(self.force_field, vec)
 
     def _energy_kernel(self) -> Callable:
         """Return this session's compiled ``energy_fn(params, coords)`` kernel.
@@ -1279,17 +1155,10 @@ def _energy_kernel(self) -> Callable:
         conformers can share a compiled topology kernel without exposing the
         private native state as a public handle.
         """
+        assert self._state._energy_fn is not None
         return self._state._energy_fn
 
-    def _energy(self, request: EnergyRequest) -> EnergyResult:  # type: ignore[override]
-        ff = self._ff_for(request.parameters)
-        try:
-            value = self._backend._evaluate_energy(self._state, ff)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX energy evaluation failed: {exc}") from exc
-        return EnergyResult(energy=float(value), unit=EnergyUnit.KCAL_PER_MOL, provenance=self._info.provenance)
-
-    def _minimize(self, request: MinimizationRequest) -> GeometryResult:  # type: ignore[override]
+    def _minimize(self, request: MinimizationRequest) -> GeometryResult:
         ff = self._ff_for(request.parameters)
         max_iterations = request.max_iterations if request.max_iterations is not None else 200
         tolerance = request.tolerance if request.tolerance is not None else 1e-5
@@ -1308,46 +1177,10 @@ def _minimize(self, request: MinimizationRequest) -> GeometryResult:  # type: ig
             symbols=tuple(atoms),
             coordinates=readonly_array(coords),
             coordinate_unit=LengthUnit.ANGSTROM,
-            provenance=self._info.provenance,
+            provenance=self._provenance,
         )
 
-    def _hessian(self, request: HessianRequest) -> HessianResult:  # type: ignore[override]
-        ff = self._ff_for(request.parameters)
-        try:
-            hess = self._backend._evaluate_hessian(self._state, ff)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX Hessian evaluation failed: {exc}") from exc
-        return HessianResult(
-            hessian=readonly_array(hess), unit=HessianUnit.HARTREE_PER_BOHR2, provenance=self._info.provenance
-        )
-
-    def _frequencies(self, request: FrequencyRequest) -> FrequencyResult:  # type: ignore[override]
-        from q2mm.models.hessian import hessian_to_frequencies
-
-        ff = self._ff_for(request.parameters)
-        try:
-            hess_au = self._backend._evaluate_hessian(self._state, ff)
-            freqs = hessian_to_frequencies(hess_au, list(self.molecule.symbols), on_error=request.on_error)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX frequency evaluation failed: {exc}") from exc
-        return FrequencyResult(
-            frequencies=readonly_array(freqs), unit=FrequencyUnit.INVERSE_CM, provenance=self._info.provenance
-        )
-
-    def _parameter_gradient(self, request: ParameterGradientRequest) -> ParameterGradientResult:  # type: ignore[override]
-        ff = self._ff_for(request.parameters)
-        try:
-            energy, grad = self._backend._evaluate_param_grad(self._state, ff)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX parameter-gradient evaluation failed: {exc}") from exc
-        return ParameterGradientResult(
-            energy=float(energy),
-            gradient=readonly_array(grad),
-            unit=EnergyUnit.KCAL_PER_MOL,
-            provenance=self._info.provenance,
-        )
-
-    def _hessian_parameter_jacobian(self, request: HessianJacobianRequest) -> HessianJacobianResult:  # type: ignore[override]
+    def _hessian_parameter_jacobian(self, request: HessianJacobianRequest) -> HessianJacobianResult:
         ff = self._ff_for(request.parameters)
         try:
             hess, jac = self._backend._evaluate_hessian_param_jac(self._state, ff)
@@ -1357,17 +1190,7 @@ def _hessian_parameter_jacobian(self, request: HessianJacobianRequest) -> Hessia
             hessian=readonly_array(hess),
             jacobian=readonly_array(jac),
             unit=HessianUnit.HARTREE_PER_BOHR2,
-            provenance=self._info.provenance,
-        )
-
-    def _batched_energy(self, request: BatchedEnergyRequest) -> BatchedEnergyResult:  # type: ignore[override]
-        mat = self._validate_matrix(request.parameter_matrix)
-        try:
-            energies = self._backend._evaluate_batched_energy(self._state, self.force_field, mat)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX batched-energy evaluation failed: {exc}") from exc
-        return BatchedEnergyResult(
-            energies=readonly_array(energies), unit=EnergyUnit.KCAL_PER_MOL, provenance=self._info.provenance
+            provenance=self._provenance,
         )
 
 
@@ -1542,100 +1365,65 @@ def _compile_energy_fn(state: _JaxState, forcefield: ForceField) -> Callable:
     _dipole_bond_indices = jnp.array(_dipole_bond_indices) if has_dipoles else None
     _dipole_pair_indices = jnp.array(_dipole_pair_indices) if has_dipoles else None
 
-    if use_mm3:
-
-        @jax.jit
-        def energy_fn(params: jnp.ndarray, coords: jnp.ndarray) -> jnp.ndarray:
-            """Compute total MM3 energy from parameter and coordinate arrays."""
-            E = jnp.float64(0.0)
-
-            if has_bonds:
-                bond_params = params[bond_offset : bond_offset + 2 * n_bt].reshape(n_bt, 2)
-                k = bond_params[_bond_map, 0] * _BOND_K_CONV
-                r0 = bond_params[_bond_map, 1]
-                E = E + _mm3_bond_energy(k, r0, coords, _bond_indices)
-
-            if has_angles:
-                angle_params = params[angle_offset : angle_offset + 2 * n_at].reshape(n_at, 2)
-                k = angle_params[_angle_map, 0] * _ANGLE_K_CONV
-                theta0_deg = angle_params[_angle_map, 1]
-                theta0 = theta0_deg * (jnp.pi / 180.0)
-                E = E + _mm3_angle_energy(k, theta0, coords, _angle_indices)
-
-            if has_torsions:
-                torsion_k = params[torsion_offset : torsion_offset + n_tt]
-                k = torsion_k[_torsion_map]
-                E = E + _torsion_energy(k, _torsion_n, _torsion_gamma, coords, _torsion_indices)
-
-            if has_vdw:
-                vdw_params = params[vdw_offset : vdw_offset + 2 * n_vt].reshape(n_vt, 2)
-                per_atom_radius = vdw_params[_atom_vdw_map, 0]
-                per_atom_epsilon = vdw_params[_atom_vdw_map, 1]
+    bond_energy_fn = _mm3_bond_energy if use_mm3 else _harmonic_bond_energy
+    angle_energy_fn = _mm3_angle_energy if use_mm3 else _harmonic_angle_energy
+
+    @jax.jit
+    def energy_fn(params: JaxArray, coords: JaxArray) -> JaxArray:
+        """Compute total MM energy from parameter and coordinate arrays."""
+        E = jnp.float64(0.0)
+
+        if has_bonds:
+            bond_params = params[bond_offset : bond_offset + 2 * n_bt].reshape(n_bt, 2)
+            k = bond_params[_bond_map, 0] * _BOND_K_CONV
+            r0 = bond_params[_bond_map, 1]
+            E = E + bond_energy_fn(k, r0, coords, _bond_indices)
+
+        if has_angles:
+            angle_params = params[angle_offset : angle_offset + 2 * n_at].reshape(n_at, 2)
+            k = angle_params[_angle_map, 0] * _ANGLE_K_CONV
+            theta0_deg = angle_params[_angle_map, 1]
+            theta0 = theta0_deg * (jnp.pi / 180.0)
+            E = E + angle_energy_fn(k, theta0, coords, _angle_indices)
+
+        if has_torsions:
+            torsion_k = params[torsion_offset : torsion_offset + n_tt]
+            k = torsion_k[_torsion_map]
+            E = E + _torsion_energy(k, _torsion_n, _torsion_gamma, coords, _torsion_indices)
+
+        if has_vdw:
+            vdw_params = params[vdw_offset : vdw_offset + 2 * n_vt].reshape(n_vt, 2)
+            per_atom_radius = vdw_params[_atom_vdw_map, 0]
+            per_atom_epsilon = vdw_params[_atom_vdw_map, 1]
+            if use_mm3:
                 E = E + _mm3_vdw_energy(per_atom_radius, per_atom_epsilon, coords, _vdw_pairs)
-
-            if has_ub:
-                ub_params = params[ub_offset : ub_offset + 2 * n_ubt].reshape(n_ubt, 2)
-                k = ub_params[_ub_map, 0]
-                r0 = ub_params[_ub_map, 1]
-                E = E + _harmonic_bond_energy(k, r0, coords, _ub_indices)
-
-            if has_sb:
-                sb_k = params[sb_offset : sb_offset + n_sbt]
-                k = sb_k[_sb_map]
-                # Read equilibrium values from the runtime param vector
-                bond_params_sb = params[bond_offset : bond_offset + 2 * n_bt].reshape(n_bt, 2)
-                angle_params_sb = params[angle_offset : angle_offset + 2 * n_at].reshape(n_at, 2)
-                r0_ij = bond_params_sb[_sb_bond_ij_idx, 1]
-                r0_jk = bond_params_sb[_sb_bond_jk_idx, 1]
-                theta0 = angle_params_sb[_sb_angle_idx, 1] * (jnp.pi / 180.0)
-                E = E + _mm3_stretch_bend_energy(k, r0_ij, r0_jk, theta0, coords, _sb_angle_indices)
-
-            if has_dipoles:
-                E = E + _mm3_dipole_energy(_dipole_moments, coords, _dipole_bond_indices, _dipole_pair_indices)
-
-            return E
-
-    else:
-
-        @jax.jit
-        def energy_fn(params: jnp.ndarray, coords: jnp.ndarray) -> jnp.ndarray:
-            """Compute total harmonic energy from parameter and coordinate arrays."""
-            E = jnp.float64(0.0)
-
-            if has_bonds:
-                bond_params = params[bond_offset : bond_offset + 2 * n_bt].reshape(n_bt, 2)
-                k = bond_params[_bond_map, 0] * _BOND_K_CONV
-                r0 = bond_params[_bond_map, 1]
-                E = E + _harmonic_bond_energy(k, r0, coords, _bond_indices)
-
-            if has_angles:
-                angle_params = params[angle_offset : angle_offset + 2 * n_at].reshape(n_at, 2)
-                k = angle_params[_angle_map, 0] * _ANGLE_K_CONV
-                theta0_deg = angle_params[_angle_map, 1]
-                theta0 = theta0_deg * (jnp.pi / 180.0)
-                E = E + _harmonic_angle_energy(k, theta0, coords, _angle_indices)
-
-            if has_torsions:
-                torsion_k = params[torsion_offset : torsion_offset + n_tt]
-                k = torsion_k[_torsion_map]
-                E = E + _torsion_energy(k, _torsion_n, _torsion_gamma, coords, _torsion_indices)
-
-            if has_vdw:
-                vdw_params = params[vdw_offset : vdw_offset + 2 * n_vt].reshape(n_vt, 2)
-                per_atom_radius = vdw_params[_atom_vdw_map, 0]
-                per_atom_epsilon = vdw_params[_atom_vdw_map, 1]
+            else:
                 per_atom_sigma = per_atom_radius * 2.0 / (2.0 ** (1.0 / 6.0))
                 E = E + _lj_12_6_energy(per_atom_sigma, per_atom_epsilon, coords, _vdw_pairs)
 
-            if has_ub:
-                ub_params = params[ub_offset : ub_offset + 2 * n_ubt].reshape(n_ubt, 2)
-                k = ub_params[_ub_map, 0]
-                r0 = ub_params[_ub_map, 1]
-                E = E + _harmonic_bond_energy(k, r0, coords, _ub_indices)
-
-            return E
-
-    return energy_fn
+        if has_ub:
+            ub_params = params[ub_offset : ub_offset + 2 * n_ubt].reshape(n_ubt, 2)
+            k = ub_params[_ub_map, 0]
+            r0 = ub_params[_ub_map, 1]
+            E = E + _harmonic_bond_energy(k, r0, coords, _ub_indices)
+
+        if use_mm3 and has_sb:
+            sb_k = params[sb_offset : sb_offset + n_sbt]
+            k = sb_k[_sb_map]
+            bond_params_sb = params[bond_offset : bond_offset + 2 * n_bt].reshape(n_bt, 2)
+            angle_params_sb = params[angle_offset : angle_offset + 2 * n_at].reshape(n_at, 2)
+            r0_ij = bond_params_sb[_sb_bond_ij_idx, 1]
+            r0_jk = bond_params_sb[_sb_bond_jk_idx, 1]
+            theta0 = angle_params_sb[_sb_angle_idx, 1] * (jnp.pi / 180.0)
+            E = E + _mm3_stretch_bend_energy(k, r0_ij, r0_jk, theta0, coords, _sb_angle_indices)
+
+        if use_mm3 and has_dipoles:
+            E = E + _mm3_dipole_energy(_dipole_moments, coords, _dipole_bond_indices, _dipole_pair_indices)
+
+        return E
+
+    compiled_fn: Callable = energy_fn
+    return compiled_fn
 
 
 # ---------------------------------------------------------------------------
diff --git a/q2mm/backends/mm/jax_md_engine.py b/q2mm/backends/mm/jax_md_engine.py
index aae64e75..32f694d6 100644
--- a/q2mm/backends/mm/jax_md_engine.py
+++ b/q2mm/backends/mm/jax_md_engine.py
@@ -48,67 +48,59 @@
 import logging
 from collections.abc import Callable
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
 logger = logging.getLogger(__name__)
 
 from q2mm.backends.contracts import (
-    AbstractPreparedBackend,
     BackendInfo,
     BackendProvenance,
     BackendRole,
     BackendUnavailableError,
-    BatchedEnergyRequest,
-    BatchedEnergyResult,
     Capability,
-    EnergyRequest,
-    EnergyResult,
     EnergyUnit,
     EvaluationError,
-    FrequencyRequest,
-    FrequencyResult,
-    FrequencyUnit,
     GeometryResult,
-    HessianRequest,
-    HessianResult,
-    HessianUnit,
     LengthUnit,
     MinimizationRequest,
-    ParameterGradientRequest,
-    ParameterGradientResult,
     PreparationError,
     PreparationRequest,
     readonly_array,
 )
-from q2mm.backends.mm._term_support import _validate_term_support
+from q2mm.backends.mm._term_support import _require_mm_force_field
 from q2mm.models.units import KCALMOLA2_TO_HESSIAN_AU
 from q2mm.models.forcefield import ForceField
 from q2mm.models.molecule import Molecule
 
 from q2mm.backends.mm._jax_common import (
+    _PreparedJaxBase,
     ensure_jax as _ensure_jax_common,
     layout_block_offsets,
-    match_angle as _match_angle,
-    match_bond as _match_bond,
-    match_vdw as _match_vdw,
+    match_mm_topology,
     params_and_coords as _params_and_coords_impl,
 )
 
+if TYPE_CHECKING:
+    from q2mm.models.parameters import ParameterLayout
+
+type JaxArray = Any
+
 # Lazy placeholders — populated by _ensure_jax_md().
-jax = None
-jnp = None
+jax: Any = None
+jnp: Any = None
 
 # Cheap availability check — does NOT import jax-md or initialize CUDA.
 _HAS_JAX_MD: bool = importlib.util.find_spec("jax_md") is not None
 
 # Lazy-loaded jax-md symbols (populated by _ensure_jax_md).
-NonbondedOptions = None
-CoulombHandler = None
-CutoffCoulomb = None
-oplsaa_energy = None
-create_parameters = None
-create_topology = None
+NonbondedOptions: Any = None
+CoulombHandler: Any = None
+CutoffCoulomb: Any = None
+oplsaa_energy: Any = None
+create_parameters: Any = None
+create_topology: Any = None
 
 # Hessian unit conversion imported from q2mm.models.units
 
@@ -200,11 +192,12 @@ class _JaxMdState:
     # Compiled functions (lazy, JIT-compiled)
     _energy_fn: Callable | None = field(default=None, repr=False)
     _scalar_energy_fn: Callable | None = field(default=None, repr=False)
+    _batched_energy_fn: Callable | None = field(default=None, repr=False)
     _grad_fn: Callable | None = field(default=None, repr=False)
     _coord_hess_fn: Callable | None = field(default=None, repr=False)
     # jax-md neighbor list function
-    _neighbor_fn: object | None = field(default=None, repr=False)
-    _nlist: object | None = field(default=None, repr=False)
+    _neighbor_fn: Any = field(default=None, repr=False)
+    _nlist: Any = field(default=None, repr=False)
 
 
 # ---------------------------------------------------------------------------
@@ -216,10 +209,10 @@ def _build_jaxmd_params_fn(
     state: _JaxMdState,
     forcefield: ForceField,
 ) -> Callable[
-    [jnp.ndarray],
+    [JaxArray],
     tuple[
-        tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray],
-        tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
+        tuple[JaxArray, JaxArray, JaxArray, JaxArray, JaxArray],
+        tuple[JaxArray, JaxArray, JaxArray],
     ],
 ]:
     """Create a function that maps a flat param vector to jax-md Parameters.
@@ -263,10 +256,10 @@ def _build_jaxmd_params_fn(
     n_torsions = len(state.torsion_param_map)
 
     def build_params(
-        param_vector: jnp.ndarray,
+        param_vector: JaxArray,
     ) -> tuple[
-        tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray],
-        tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
+        tuple[JaxArray, JaxArray, JaxArray, JaxArray, JaxArray],
+        tuple[JaxArray, JaxArray, JaxArray],
     ]:
         """Unpack a flat parameter vector into per-term JAX-MD arrays.
 
@@ -333,8 +326,8 @@ def build_params(
 
 
 def _compile_energy_fn(
-    state: _JaxMdState, forcefield: ForceField, coulomb: object, nb_options: object
-) -> tuple[Callable, Callable, object, object]:
+    state: _JaxMdState, forcefield: ForceField, coulomb: Any, nb_options: Any
+) -> tuple[Callable, Callable, Any, Any]:
     """Compile the jax-md energy function for a specific topology.
 
     Args:
@@ -451,7 +444,7 @@ def _compile_energy_fn(
     has_torsions = torsion_arr.shape[0] > 0
     has_vdw = state.n_vdw_types > 0
 
-    def energy_fn(param_vector: jnp.ndarray, coords: jnp.ndarray, nlist_: object) -> dict[str, jnp.ndarray]:
+    def energy_fn(param_vector: JaxArray, coords: JaxArray, nlist_: Any) -> dict[str, JaxArray]:
         """Compute total energy given param vector and coordinates.
 
         Returns dict with per-term breakdown.
@@ -487,8 +480,9 @@ def energy_fn(param_vector: jnp.ndarray, coords: jnp.ndarray, nlist_: object) ->
         # Torsion energy: E = k * (1 + cos(n*phi - gamma))
         if has_torsions:
             idx = _torsion_indices
+            assert idx is not None
 
-            def compute_dihedral(p0: jnp.ndarray, p1: jnp.ndarray, p2: jnp.ndarray, p3: jnp.ndarray) -> jnp.ndarray:
+            def compute_dihedral(p0: JaxArray, p1: JaxArray, p2: JaxArray, p3: JaxArray) -> JaxArray:
                 """Compute signed dihedral angle (radians) for four points."""
                 b0 = displacement_fn(p1, p0)
                 b1 = displacement_fn(p2, p1)
@@ -568,7 +562,7 @@ def compute_dihedral(p0: jnp.ndarray, p1: jnp.ndarray, p2: jnp.ndarray, p3: jnp.
             "total": E_total,
         }
 
-    def scalar_energy_fn(param_vector: jnp.ndarray, coords: jnp.ndarray, nlist_: object) -> jnp.ndarray:
+    def scalar_energy_fn(param_vector: JaxArray, coords: JaxArray, nlist_: object) -> JaxArray:
         """Scalar total energy for gradient/hessian computation."""
         return energy_fn(param_vector, coords, nlist_)["total"]
 
@@ -603,8 +597,8 @@ class JaxMdBackend:
     def __init__(
         self,
         box: tuple[float, float, float] = (100.0, 100.0, 100.0),
-        coulomb: CoulombHandler | None = None,
-        nb_options: NonbondedOptions | None = None,
+        coulomb: Any = None,
+        nb_options: Any = None,
     ) -> None:
         if not _HAS_JAX_MD:
             raise BackendUnavailableError("jax-md is not installed. Install with `pip install jax-md`.")
@@ -663,22 +657,16 @@ def prepare(self, request: PreparationRequest) -> PreparedJaxMd:
         """
         from q2mm.models.parameters import ParameterLayout
 
-        if request.force_field is None:
-            raise PreparationError("JAX-MD requires a base ForceField in the PreparationRequest.")
         info = self.info
-        form = request.force_field.functional_form.value
-        if not info.supports_form(form):
-            raise PreparationError(
-                f"JAX-MD does not support functional form {form!r}. Supported: {sorted(info.functional_forms)}"
-            )
-        _validate_term_support(
-            request.force_field,
+        force_field = _require_mm_force_field(
+            request,
+            info,
             backend="JAX-MD",
             unsupported=frozenset({"Urey-Bradley", "CMAP", "improper torsions", "bond dipoles", "vdW reduction"}),
         )
-        layout = ParameterLayout.from_force_field(request.force_field)
+        layout = ParameterLayout.from_force_field(force_field)
         try:
-            state = self._build_state(request.molecule, request.force_field)
+            state = self._build_state(request.molecule, force_field)
         except PreparationError:
             raise
         except Exception as exc:  # noqa: BLE001
@@ -688,7 +676,7 @@ def prepare(self, request: PreparationRequest) -> PreparedJaxMd:
             info=info,
             case_id=request.case_id,
             molecule=request.molecule,
-            force_field=request.force_field,
+            force_field=force_field,
             layout=layout,
             state=state,
         )
@@ -713,89 +701,18 @@ def _build_state(self, molecule: Molecule, forcefield: ForceField) -> _JaxMdStat
                 "Use the JAX backend for force fields with stretch-bend parameters."
             )
 
-        # Match bonds
-        bond_atom_indices = []
-        bond_param_map = []
-        for bond in molecule.bonds:
-            idx, param = _match_bond(
-                forcefield,
-                bond.elements,
-                env_id=bond.env_id,
-                ff_row=bond.ff_row,
-                bond_order=getattr(bond, "bond_order", ""),
-                bond_length=bond.length,
-            )
-            if param is not None:
-                bond_atom_indices.append((bond.atom_i, bond.atom_j))
-                bond_param_map.append(idx)
-
-        # Match angles
-        angle_atom_indices = []
-        angle_param_map = []
-        for angle in molecule.angles:
-            idx, param = _match_angle(forcefield, angle.elements, env_id=angle.env_id, ff_row=angle.ff_row)
-            if param is not None:
-                angle_atom_indices.append((angle.atom_i, angle.atom_j, angle.atom_k))
-                angle_param_map.append(idx)
-
-        # Match torsions — each detected torsion may match multiple FF
-        # entries (one per periodicity component)
-        torsion_atom_indices: list[tuple[int, int, int, int]] = []
-        torsion_param_map: list[int] = []
-        torsion_param_index = {id(p): i for i, p in enumerate(forcefield.torsions)}
-        for torsion in molecule.torsions:
-            matches = forcefield.match_torsion(
-                torsion.element_quad, env_id=torsion.env_id, ff_row=torsion.ff_row, is_improper=False
-            )
-            for param in matches:
-                j_ff = torsion_param_index[id(param)]
-                torsion_atom_indices.append((torsion.atom_i, torsion.atom_j, torsion.atom_k, torsion.atom_l))
-                torsion_param_map.append(j_ff)
-
-        # Match vdW
-        atom_vdw_map = []
-        excluded_types = {value.casefold() for value in forcefield.nonbonded_excluded_atom_types}
-        for symbol, atom_type in zip(molecule.symbols, molecule.atom_types, strict=False):
-            if symbol.casefold() in excluded_types or atom_type.casefold() in excluded_types:
-                atom_vdw_map.append(-1)
-                continue
-            idx, param = _match_vdw(forcefield, atom_type=atom_type, element=symbol)
-            atom_vdw_map.append(idx if idx is not None else -1)
-
-        # Validate vdW
-        explicitly_excluded = {
-            index
-            for index, (symbol, atom_type) in enumerate(zip(molecule.symbols, molecule.atom_types, strict=True))
-            if symbol.casefold() in excluded_types or atom_type.casefold() in excluded_types
-        }
-        unmatched = [i for i, idx in enumerate(atom_vdw_map) if idx == -1 and i not in explicitly_excluded]
-        if getattr(forcefield, "vdws", None) and unmatched:
-            raise ValueError(
-                f"Unmatched vdW parameters for atoms at indices {unmatched}. "
-                "Ensure all atom types/elements have corresponding vdW "
-                "parameters in the force field."
-            )
-
-        bond_indices_arr = (
-            np.array(bond_atom_indices, dtype=np.int32) if bond_atom_indices else np.empty((0, 2), dtype=np.int32)
-        )
-        angle_indices_arr = (
-            np.array(angle_atom_indices, dtype=np.int32) if angle_atom_indices else np.empty((0, 3), dtype=np.int32)
-        )
-        torsion_indices_arr = (
-            np.array(torsion_atom_indices, dtype=np.int32) if torsion_atom_indices else np.empty((0, 4), dtype=np.int32)
-        )
+        topo = match_mm_topology(molecule, forcefield, include_impropers=False, mask_excluded_vdw=True)
 
         state = _JaxMdState(
             molecule=molecule,
             box=self._box.copy(),
-            bond_indices=bond_indices_arr,
-            angle_indices=angle_indices_arr,
-            torsion_indices=torsion_indices_arr,
-            bond_param_map=np.array(bond_param_map, dtype=np.int32),
-            angle_param_map=np.array(angle_param_map, dtype=np.int32),
-            torsion_param_map=np.array(torsion_param_map, dtype=np.int32),
-            atom_vdw_map=np.array(atom_vdw_map, dtype=np.int32),
+            bond_indices=topo.bond_indices,
+            angle_indices=topo.angle_indices,
+            torsion_indices=topo.torsion_indices,
+            bond_param_map=topo.bond_param_map,
+            angle_param_map=topo.angle_param_map,
+            torsion_param_map=topo.torsion_param_map,
+            atom_vdw_map=topo.atom_vdw_map,
             n_bond_types=len(forcefield.bonds),
             n_angle_types=len(forcefield.angles),
             n_torsion_types=len(forcefield.torsions),
@@ -818,19 +735,21 @@ def _build_state(self, molecule: Molecule, forcefield: ForceField) -> _JaxMdStat
 
         return state
 
-    def _params_and_coords(self, state: _JaxMdState, forcefield: ForceField) -> tuple[jnp.ndarray, jnp.ndarray]:
+    def _params_and_coords(self, state: _JaxMdState, forcefield: ForceField) -> tuple[JaxArray, JaxArray]:
         """Extract JAX arrays from force field and molecule."""
         return _params_and_coords_impl(state.molecule.geometry, forcefield)
 
     def _evaluate_energy(self, state: _JaxMdState, forcefield: ForceField) -> float:
         """Total energy in kcal/mol for a prepared native state."""
         params, coords = self._params_and_coords(state, forcefield)
+        assert state._energy_fn is not None
         result = state._energy_fn(params, coords, state._nlist)
         return float(result["total"])
 
     def _evaluate_energy_breakdown(self, state: _JaxMdState, forcefield: ForceField) -> dict[str, float]:
         """Per-term energy breakdown (kcal/mol) for a prepared native state."""
         params, coords = self._params_and_coords(state, forcefield)
+        assert state._energy_fn is not None
         result = state._energy_fn(params, coords, state._nlist)
         return {k: float(v) for k, v in result.items()}
 
@@ -853,7 +772,7 @@ def _evaluate_batched_energy(
         batch_params = jnp.array(param_matrix, dtype=jnp.float64)
         nlist = state._nlist
 
-        if not hasattr(state, "_batched_energy_fn") or state._batched_energy_fn is None:
+        if state._batched_energy_fn is None:
             state._batched_energy_fn = jax.jit(jax.vmap(state._scalar_energy_fn, in_axes=(0, None, None)))
 
         return np.asarray(state._batched_energy_fn(batch_params, coords, nlist))
@@ -864,9 +783,11 @@ def _evaluate_hessian(self, state: _JaxMdState, forcefield: ForceField) -> np.nd
 
         if state._coord_hess_fn is None:
             nlist = state._nlist
+            scalar_fn = state._scalar_energy_fn
+            assert scalar_fn is not None
 
-            def _energy_of_flat_coords(flat_coords: jnp.ndarray, params_: jnp.ndarray) -> jnp.ndarray:
-                return state._scalar_energy_fn(params_, flat_coords.reshape(-1, 3), nlist)
+            def _energy_of_flat_coords(flat_coords: JaxArray, params_: JaxArray) -> JaxArray:
+                return scalar_fn(params_, flat_coords.reshape(-1, 3), nlist)
 
             state._coord_hess_fn = jax.jit(jax.hessian(_energy_of_flat_coords, argnums=0))
 
@@ -885,13 +806,13 @@ def _evaluate_minimize(self, state: _JaxMdState, forcefield: ForceField, *, max_
 
         params, coords = self._params_and_coords(state, forcefield)
         neighbor_fn = state._neighbor_fn
-
         scalar_fn = state._scalar_energy_fn
+        assert scalar_fn is not None
 
         # Mutable container for neighbor list — updated when atoms move
         nlist_ref = [state._nlist]
 
-        def _update_nlist(new_coords: jnp.ndarray) -> None:
+        def _update_nlist(new_coords: JaxArray) -> None:
             """Re-allocate neighbor list for new coordinates."""
             nlist_ref[0] = neighbor_fn.allocate(new_coords)
 
@@ -924,13 +845,15 @@ def gradient(x: np.ndarray) -> np.ndarray:
         return opt_energy, list(state.molecule.symbols), opt_coords
 
 
-class PreparedJaxMd(AbstractPreparedBackend):
+class PreparedJaxMd(_PreparedJaxBase):
     """Prepared JAX-MD session for a single training case.
 
     Owns the molecule, base force field, parameter layout, and one compiled
     private :class:`_JaxMdState` reused across parameter vectors.
     """
 
+    _backend_label = "JAX-MD"
+
     def __init__(
         self,
         *,
@@ -939,32 +862,20 @@ def __init__(
         case_id: str,
         molecule: Molecule,
         force_field: ForceField,
-        layout: object,
+        layout: ParameterLayout,
         state: _JaxMdState,
     ) -> None:
         super().__init__(
+            backend=backend,
             info=info,
             case_id=case_id,
             molecule=molecule,
             force_field=force_field,
-            layout=layout,  # type: ignore[arg-type]
+            layout=layout,
+            state=state,
         )
-        self._backend = backend
-        self._state = state
-
-    def _ff_for(self, parameters: np.ndarray) -> ForceField:
-        vec = self._validate_vector(parameters)
-        return self.layout.replace(self.force_field, vec)
-
-    def _energy(self, request: EnergyRequest) -> EnergyResult:  # type: ignore[override]
-        ff = self._ff_for(request.parameters)
-        try:
-            value = self._backend._evaluate_energy(self._state, ff)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX-MD energy evaluation failed: {exc}") from exc
-        return EnergyResult(energy=float(value), unit=EnergyUnit.KCAL_PER_MOL, provenance=self._info.provenance)
 
-    def _minimize(self, request: MinimizationRequest) -> GeometryResult:  # type: ignore[override]
+    def _minimize(self, request: MinimizationRequest) -> GeometryResult:
         ff = self._ff_for(request.parameters)
         if request.tolerance is not None:
             raise EvaluationError(
@@ -982,53 +893,7 @@ def _minimize(self, request: MinimizationRequest) -> GeometryResult:  # type: ig
             symbols=tuple(atoms),
             coordinates=readonly_array(coords),
             coordinate_unit=LengthUnit.ANGSTROM,
-            provenance=self._info.provenance,
-        )
-
-    def _hessian(self, request: HessianRequest) -> HessianResult:  # type: ignore[override]
-        ff = self._ff_for(request.parameters)
-        try:
-            hess = self._backend._evaluate_hessian(self._state, ff)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX-MD Hessian evaluation failed: {exc}") from exc
-        return HessianResult(
-            hessian=readonly_array(hess), unit=HessianUnit.HARTREE_PER_BOHR2, provenance=self._info.provenance
-        )
-
-    def _frequencies(self, request: FrequencyRequest) -> FrequencyResult:  # type: ignore[override]
-        from q2mm.models.hessian import hessian_to_frequencies
-
-        ff = self._ff_for(request.parameters)
-        try:
-            hess_au = self._backend._evaluate_hessian(self._state, ff)
-            freqs = hessian_to_frequencies(hess_au, list(self.molecule.symbols), on_error=request.on_error)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX-MD frequency evaluation failed: {exc}") from exc
-        return FrequencyResult(
-            frequencies=readonly_array(freqs), unit=FrequencyUnit.INVERSE_CM, provenance=self._info.provenance
-        )
-
-    def _parameter_gradient(self, request: ParameterGradientRequest) -> ParameterGradientResult:  # type: ignore[override]
-        ff = self._ff_for(request.parameters)
-        try:
-            energy, grad = self._backend._evaluate_param_grad(self._state, ff)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX-MD parameter-gradient evaluation failed: {exc}") from exc
-        return ParameterGradientResult(
-            energy=float(energy),
-            gradient=readonly_array(grad),
-            unit=EnergyUnit.KCAL_PER_MOL,
-            provenance=self._info.provenance,
-        )
-
-    def _batched_energy(self, request: BatchedEnergyRequest) -> BatchedEnergyResult:  # type: ignore[override]
-        mat = self._validate_matrix(request.parameter_matrix)
-        try:
-            energies = self._backend._evaluate_batched_energy(self._state, self.force_field, mat)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"JAX-MD batched-energy evaluation failed: {exc}") from exc
-        return BatchedEnergyResult(
-            energies=readonly_array(energies), unit=EnergyUnit.KCAL_PER_MOL, provenance=self._info.provenance
+            provenance=self._provenance,
         )
 
 
diff --git a/q2mm/backends/mm/openmm.py b/q2mm/backends/mm/openmm.py
index 6e35493d..734bde8f 100644
--- a/q2mm/backends/mm/openmm.py
+++ b/q2mm/backends/mm/openmm.py
@@ -29,7 +29,6 @@
     EvaluationError,
     FrequencyRequest,
     FrequencyResult,
-    FrequencyUnit,
     GeometryResult,
     HessianRequest,
     HessianResult,
@@ -40,9 +39,12 @@
     ParameterGradientResult,
     PreparationError,
     PreparationRequest,
+    ReferenceEnergyRequest,
+    ReferenceFrequencyRequest,
+    ReferenceHessianRequest,
     readonly_array,
 )
-from q2mm.backends.mm._term_support import _validate_term_support
+from q2mm.backends.mm._term_support import _require_mm_force_field
 from q2mm.backends.mm._openmm_terms import (
     _AngleTerm,
     _BondTerm,
@@ -83,17 +85,18 @@
 from q2mm.models.forcefield import AngleParam, BondParam, ForceField, FunctionalForm, TorsionParam, VdwParam
 from q2mm.models.molecule import Molecule
 
+mm: Any = None
+unit: Any = None
+_OpenMMException: Any = Exception
+_HAS_OPENMM = False
+
 try:
-    from openmm import OpenMMException as _OpenMMException
-    from openmm import openmm as mm
-    from openmm import unit
+    from openmm import OpenMMException as _OpenMMException_cls, openmm as _mm_mod, unit as _unit_mod
 
+    mm, unit, _OpenMMException = _mm_mod, _unit_mod, _OpenMMException_cls
     _HAS_OPENMM = True
 except ImportError:  # pragma: no cover - exercised when OpenMM is not installed
-    mm = None
-    unit = None
-    _OpenMMException = Exception  # type: ignore[assignment,misc]
-    _HAS_OPENMM = False
+    pass
 
 
 @dataclass
@@ -125,15 +128,15 @@ class _OpenMMState:
     """
 
     molecule: Molecule
-    system: object
-    integrator: object
-    context: object
-    bond_force: object | None
-    angle_force: object | None
-    torsion_force: object | None
-    vdw_force: object | None
-    ub_force: object | None
-    cmap_force: object | None
+    system: Any
+    integrator: Any
+    context: Any
+    bond_force: Any
+    angle_force: Any
+    torsion_force: Any
+    vdw_force: Any
+    ub_force: Any
+    cmap_force: Any
     bond_terms: list[_BondTerm]
     bond_parameters: tuple[BondParam, ...]
     angle_terms: list[_AngleTerm]
@@ -164,8 +167,8 @@ class _OpenMMDiffState:
 
     """
 
-    integrator: object
-    context: object
+    integrator: Any
+    context: Any
     param_names: list[str]
     param_vector_indices: list[int]
     grad_unit_factors: list[float]
@@ -251,7 +254,7 @@ def detect_best_platform() -> str:
                     )
             return name
     # Fallback — shouldn't happen since OpenMM always has Reference
-    return mm.Platform.getPlatform(0).getName()  # pragma: no cover
+    return str(mm.Platform.getPlatform(0).getName())  # pragma: no cover
 
 
 def _collect_bond_assignments(
@@ -264,7 +267,7 @@ def _collect_bond_assignments(
     molecule traversal order.  Bonds with no matching parameter are skipped.
     """
     assignments: list[tuple[Any, BondParam]] = []
-    for bond in molecule.bonds:
+    for bond in molecule.bonds or ():
         param = forcefield.match_bond(
             bond.elements,
             env_id=bond.env_id,
@@ -306,7 +309,7 @@ def _collect_angle_assignments(
     molecule traversal order.  Angles with no matching parameter are skipped.
     """
     assignments: list[tuple[Any, AngleParam]] = []
-    for angle in molecule.angles:
+    for angle in molecule.angles or ():
         param = forcefield.match_angle(angle.elements, env_id=angle.env_id, ff_row=angle.ff_row)
         if param is not None:
             assignments.append((angle, param))
@@ -330,7 +333,7 @@ def _collect_torsion_assignments(
         is_improper: If ``True``, match improper torsions; otherwise proper.
 
     """
-    source = molecule.improper_torsions if is_improper else molecule.torsions
+    source = molecule.improper_torsions if is_improper else (molecule.torsions or ())
     assignments: list[tuple[Any, TorsionParam]] = []
     for torsion in source:
         params = forcefield.match_torsion(
@@ -359,8 +362,9 @@ def _collect_vdw_assignments(
 
     """
     assignments: list[tuple[int, str, str, VdwParam | None]] = []
+    atom_types = molecule.atom_types if molecule.atom_types is not None else molecule.symbols
     excluded_types = {value.casefold() for value in forcefield.nonbonded_excluded_atom_types}
-    for atom_index, (symbol, atom_type) in enumerate(zip(molecule.symbols, molecule.atom_types, strict=False)):
+    for atom_index, (symbol, atom_type) in enumerate(zip(molecule.symbols, atom_types, strict=False)):
         explicitly_excluded = symbol.casefold() in excluded_types or atom_type.casefold() in excluded_types
         param = None if explicitly_excluded else forcefield.match_vdw(atom_type=atom_type, element=symbol)
         if param is None and not explicitly_excluded:
@@ -387,6 +391,13 @@ def _index_by_param_id(
     return by_param
 
 
+def _exception_14_params(vdw_force: Any, p1: int, p2: int) -> tuple[Any, Any]:
+    """Compute AMBER-scaled 1-4 LJ (sigma, epsilon) from two particle indices."""
+    _, sig1, eps1 = vdw_force.getParticleParameters(p1)
+    _, sig2, eps2 = vdw_force.getParticleParameters(p2)
+    return 0.5 * (sig1 + sig2), (eps1 * eps2) ** 0.5 / 2.0
+
+
 def _build_harmonic_exclusions(
     molecule: Molecule,
     vdw_force: Any,
@@ -408,13 +419,9 @@ def _build_harmonic_exclusions(
     for p1, p2 in sorted(excluded):
         vdw_force.addException(p1, p2, 0.0, 1.0, 0.0)
 
-    SCNB = 2.0
     exceptions_14: list[_Exception14] = []
     for p1, p2 in sorted(pairs_14):
-        sig1, eps1 = vdw_force.getParticleParameters(p1)[1:]
-        sig2, eps2 = vdw_force.getParticleParameters(p2)[1:]
-        sig_14 = 0.5 * (sig1 + sig2)
-        eps_14 = (eps1 * eps2) ** 0.5 / SCNB
+        sig_14, eps_14 = _exception_14_params(vdw_force, p1, p2)
         exc_idx = vdw_force.addException(p1, p2, 0.0, sig_14, eps_14)
         exceptions_14.append(_Exception14(exception_index=exc_idx, particle_i=p1, particle_j=p2))
 
@@ -424,7 +431,7 @@ def _build_harmonic_exclusions(
 def _bond_graph_pairs(molecule: Molecule) -> tuple[set[tuple[int, int]], set[tuple[int, int]]]:
     """Return 1-2/1-3 exclusions (shortest bond-path distances 1/2) and 1-4 pairs (distance 3)."""
     neighbors: dict[int, set[int]] = {}
-    for bond in molecule.bonds:
+    for bond in molecule.bonds or ():
         neighbors.setdefault(bond.atom_i, set()).add(bond.atom_j)
         neighbors.setdefault(bond.atom_j, set()).add(bond.atom_i)
     excluded: set[tuple[int, int]] = set()
@@ -487,10 +494,10 @@ def _find_dihedral_atoms(
     """
     results: list[tuple[int, int, int, int]] = []
     t1, t2, t3, t4 = atom_types
-    atom_types_array = molecule.atom_types
+    atom_types_array = molecule.atom_types if molecule.atom_types is not None else molecule.symbols
 
     # Use pre-enumerated torsions (bonded 1-4 paths) when available
-    for torsion in molecule.torsions:
+    for torsion in molecule.torsions or ():
         a, b, c, d = torsion.atom_i, torsion.atom_j, torsion.atom_k, torsion.atom_l
         if (
             atom_types_array[a] == t1
@@ -581,6 +588,68 @@ def _build_cmap_force(molecule: Molecule, forcefield: ForceField) -> tuple[objec
     return cmap_force, cmap_terms
 
 
+def _build_vdw_force(
+    molecule: Molecule,
+    forcefield: ForceField,
+    *,
+    use_harmonic: bool,
+) -> tuple[Any, list[_VdwTerm], list[_Exception14]]:
+    """Build the OpenMM vdW force (NonbondedForce or CustomNonbondedForce) and particle records."""
+    if use_harmonic:
+        vdw_force: Any = mm.NonbondedForce()
+        vdw_force.setNonbondedMethod(mm.NonbondedForce.NoCutoff)
+    else:
+        # MM3 Buckingham exp-6 with short-range repulsive wall.
+        # Below r < 0.34·rv the attractive r^-6 dominates the exponential,
+        # causing divergence to -∞.  Switch to a hard repulsive form
+        # at that boundary using step() to prevent collapse.
+        vdw_force = mm.CustomNonbondedForce(
+            f"step(r - rc) * epsilon*(-{MM3_VDW_C}*(rv/r)^6 + {MM3_VDW_A}*exp(-{MM3_VDW_B}*r/rv))"
+            f" + step(rc - r) * epsilon*{MM3_VDW_A}*exp(-{MM3_VDW_B}*rc/rv) * (rc/r)^12;"
+            "rc=0.34*rv;"
+            "rv=radius1+radius2;"
+            "epsilon=sqrt(epsilon1*epsilon2)"
+        )
+        vdw_force.addPerParticleParameter("radius")
+        vdw_force.addPerParticleParameter("epsilon")
+        vdw_force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
+
+    vdw_terms: list[_VdwTerm] = []
+    exceptions_14: list[_Exception14] = []
+    if forcefield.vdws:
+        vdw_assignments = _collect_vdw_assignments(molecule, forcefield)
+        for atom_index, symbol, atom_type, vdw_param in vdw_assignments:
+            if vdw_param is None:
+                if use_harmonic:
+                    vdw_force.addParticle(0.0, 0.0, 0.0)
+                else:
+                    vdw_force.addParticle([0.0, 0.0])
+                continue
+            if use_harmonic:
+                vdw_force.addParticle(0.0, _vdw_sigma_nm(vdw_param.radius), _vdw_epsilon_to_openmm(vdw_param.epsilon))
+            else:
+                vdw_force.addParticle(
+                    [_vdw_radius_to_openmm(vdw_param.radius), _vdw_epsilon_to_openmm(vdw_param.epsilon)]
+                )
+            vdw_terms.append(
+                _VdwTerm(
+                    particle_index=atom_index,
+                    atom_type=atom_type,
+                    element=symbol,
+                    ff_row=vdw_param.ff_row,
+                )
+            )
+        if use_harmonic:
+            exceptions_14 = _build_harmonic_exclusions(molecule, vdw_force)
+        else:
+            vdw_force.createExclusionsFromBonds([(bond.atom_i, bond.atom_j) for bond in molecule.bonds or ()], 2)
+            _restrict_mm3_centers(vdw_force, vdw_assignments)
+    return vdw_force, vdw_terms, exceptions_14
+
+
+_VALID_PRECISIONS = frozenset({"single", "mixed", "double"})
+
+
 class OpenMMBackend:
     """Molecular mechanics backend powered by OpenMM.
 
@@ -619,7 +688,6 @@ def __init__(
             platform_name = detect_best_platform()
         self._platform_name = platform_name
 
-        _VALID_PRECISIONS = {"single", "mixed", "double"}
         if precision is not None:
             precision = precision.strip().lower()
             if precision not in _VALID_PRECISIONS:
@@ -684,22 +752,16 @@ def prepare(self, request: PreparationRequest) -> PreparedOpenMM:
         """
         from q2mm.models.parameters import ParameterLayout
 
-        if request.force_field is None:
-            raise PreparationError("OpenMM requires a base ForceField in the PreparationRequest.")
         info = self.info
-        form = request.force_field.functional_form.value
-        if not info.supports_form(form):
-            raise PreparationError(
-                f"OpenMM does not support functional form {form!r}. Supported: {sorted(info.functional_forms)}"
-            )
-        _validate_term_support(
-            request.force_field,
+        force_field = _require_mm_force_field(
+            request,
+            info,
             backend="OpenMM",
             unsupported=frozenset({"bond dipoles", "vdW reduction", "wildcard torsions"}),
         )
-        layout = ParameterLayout.from_force_field(request.force_field)
+        layout = ParameterLayout.from_force_field(force_field)
         try:
-            state = self._build_state(request.molecule, request.force_field)
+            state = self._build_state(request.molecule, force_field)
         except BackendConfigurationError, PreparationError:
             raise
         except Exception as exc:  # noqa: BLE001
@@ -709,7 +771,7 @@ def prepare(self, request: PreparationRequest) -> PreparedOpenMM:
             info=info,
             case_id=request.case_id,
             molecule=request.molecule,
-            force_field=request.force_field,
+            force_field=force_field,
             layout=layout,
             state=state,
         )
@@ -769,7 +831,6 @@ def _build_context(self, system: Any, *, precision: str | None = None) -> tuple[
         try:
             if self._platform_name in gpu_platforms:
                 precision = precision or self._precision or "mixed"
-                _VALID_PRECISIONS = {"single", "mixed", "double"}
                 if precision not in _VALID_PRECISIONS:
                     raise BackendConfigurationError(
                         f"Invalid precision {precision!r}. Expected one of: {', '.join(sorted(_VALID_PRECISIONS))}."
@@ -871,76 +932,57 @@ def _build_state(
         # E = k * (1 + cos(n*θ − phase))
         torsion_force = mm.PeriodicTorsionForce()
 
-        if use_harmonic:
-            vdw_force = mm.NonbondedForce()
-            vdw_force.setNonbondedMethod(mm.NonbondedForce.NoCutoff)
-        else:
-            # MM3 Buckingham exp-6 with short-range repulsive wall.
-            # Below r < 0.34·rv the attractive r^-6 dominates the exponential,
-            # causing divergence to -∞.  Switch to a hard repulsive form
-            # at that boundary using step() to prevent collapse.
-            vdw_force = mm.CustomNonbondedForce(
-                f"step(r - rc) * epsilon*(-{MM3_VDW_C}*(rv/r)^6 + {MM3_VDW_A}*exp(-{MM3_VDW_B}*r/rv))"
-                f" + step(rc - r) * epsilon*{MM3_VDW_A}*exp(-{MM3_VDW_B}*rc/rv) * (rc/r)^12;"
-                "rc=0.34*rv;"
-                "rv=radius1+radius2;"
-                "epsilon=sqrt(epsilon1*epsilon2)"
-            )
-            vdw_force.addPerParticleParameter("radius")
-            vdw_force.addPerParticleParameter("epsilon")
-            vdw_force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
-
         # --- Assign bond parameters ---
         bond_assignments = (
             _collect_bond_assignments(molecule, forcefield)
             if bond_source is None
             else _bound_bond_assignments(bond_source, forcefield)
         )
-        bond_indices = {id(param): index for index, param in enumerate(forcefield.bonds)}
+        bond_indices = {id(bond_param): index for index, bond_param in enumerate(forcefield.bonds)}
         if len(bond_indices) != len(forcefield.bonds):
             raise ValueError("Cannot bind repeated bond parameter objects to a unique source index.")
         bond_terms: list[_BondTerm] = []
-        for bond, param in bond_assignments:
+        for bond, bond_param in bond_assignments:
             if use_harmonic:
                 force_index = bond_force.addBond(
                     bond.atom_i,
                     bond.atom_j,
-                    ang_to_nm(param.equilibrium),
-                    _bond_k_to_harmonic(param.force_constant),
+                    ang_to_nm(bond_param.equilibrium),
+                    _bond_k_to_harmonic(bond_param.force_constant),
                 )
             else:
                 force_index = bond_force.addBond(
                     bond.atom_i,
                     bond.atom_j,
-                    [_bond_k_to_openmm(param.force_constant), ang_to_nm(param.equilibrium)],
+                    [_bond_k_to_openmm(bond_param.force_constant), ang_to_nm(bond_param.equilibrium)],
                 )
             bond_terms.append(
                 _BondTerm(
                     force_index=force_index,
                     atom_i=bond.atom_i,
                     atom_j=bond.atom_j,
-                    parameter_index=bond_indices[id(param)],
+                    parameter_index=bond_indices[id(bond_param)],
                 )
             )
 
         # --- Assign angle parameters ---
         angle_assignments = _collect_angle_assignments(molecule, forcefield)
         angle_terms: list[_AngleTerm] = []
-        for angle, param in angle_assignments:
+        for angle, angle_param in angle_assignments:
             if use_harmonic:
                 force_index = angle_force.addAngle(
                     angle.atom_i,
                     angle.atom_j,
                     angle.atom_k,
-                    np.deg2rad(float(param.equilibrium)),
-                    _angle_k_to_harmonic(param.force_constant),
+                    np.deg2rad(float(angle_param.equilibrium)),
+                    _angle_k_to_harmonic(angle_param.force_constant),
                 )
             else:
                 force_index = angle_force.addAngle(
                     angle.atom_i,
                     angle.atom_j,
                     angle.atom_k,
-                    [_angle_k_to_openmm(param.force_constant), np.deg2rad(float(param.equilibrium))],
+                    [_angle_k_to_openmm(angle_param.force_constant), np.deg2rad(float(angle_param.equilibrium))],
                 )
             angle_terms.append(
                 _AngleTerm(
@@ -958,10 +1000,10 @@ def _build_state(
         # Reuse precomputed angle assignments to avoid rematching.
         ub_force = None
         ub_terms: list[_UBTerm] = []
-        for (angle, param), angle_term in zip(angle_assignments, angle_terms, strict=True):
-            if param.ub_force_constant is None and param.ub_equilibrium is None:
+        for (_angle, angle_param), angle_term in zip(angle_assignments, angle_terms, strict=True):
+            if angle_param.ub_force_constant is None and angle_param.ub_equilibrium is None:
                 continue
-            if param.ub_force_constant is None or param.ub_equilibrium is None:
+            if angle_param.ub_force_constant is None or angle_param.ub_equilibrium is None:
                 raise ValueError(
                     "Inconsistent Urey-Bradley parameters for angle "
                     f"{angle_term.elements} (env_id={angle_term.env_id}, ff_row={angle_term.ff_row}): "
@@ -972,8 +1014,8 @@ def _build_state(
             force_index = ub_force.addBond(
                 angle_term.atom_i,
                 angle_term.atom_k,
-                ang_to_nm(param.ub_equilibrium),
-                _bond_k_to_harmonic(param.ub_force_constant),
+                ang_to_nm(angle_param.ub_equilibrium),
+                _bond_k_to_harmonic(angle_param.ub_force_constant),
             )
             ub_terms.append(
                 _UBTerm(
@@ -992,15 +1034,15 @@ def _build_state(
 
         torsion_terms: list[_TorsionTerm] = []
         for assignments, is_improper in ((proper_assignments, False), (improper_assignments, True)):
-            for torsion, param in assignments:
+            for torsion, torsion_param in assignments:
                 force_index = torsion_force.addTorsion(
                     torsion.atom_i,
                     torsion.atom_j,
                     torsion.atom_k,
                     torsion.atom_l,
-                    param.periodicity,
-                    np.deg2rad(float(param.phase)),
-                    canonical_to_openmm_torsion_k(param.force_constant),
+                    torsion_param.periodicity,
+                    np.deg2rad(float(torsion_param.phase)),
+                    canonical_to_openmm_torsion_k(torsion_param.force_constant),
                 )
                 torsion_terms.append(
                     _TorsionTerm(
@@ -1010,41 +1052,15 @@ def _build_state(
                         atom_k=torsion.atom_k,
                         atom_l=torsion.atom_l,
                         elements=torsion.element_quad,
-                        periodicity=param.periodicity,
+                        periodicity=torsion_param.periodicity,
                         env_id=torsion.env_id,
-                        ff_row=param.ff_row,
+                        ff_row=torsion_param.ff_row,
                         is_improper=is_improper,
                     )
                 )
 
         # --- Assign vdW parameters ---
-        vdw_terms: list[_VdwTerm] = []
-        if forcefield.vdws:
-            vdw_assignments = _collect_vdw_assignments(molecule, forcefield)
-            for atom_index, symbol, atom_type, param in vdw_assignments:
-                if param is None:
-                    if use_harmonic:
-                        vdw_force.addParticle(0.0, 0.0, 0.0)
-                    else:
-                        vdw_force.addParticle([0.0, 0.0])
-                    continue
-                if use_harmonic:
-                    vdw_force.addParticle(0.0, _vdw_sigma_nm(param.radius), _vdw_epsilon_to_openmm(param.epsilon))
-                else:
-                    vdw_force.addParticle([_vdw_radius_to_openmm(param.radius), _vdw_epsilon_to_openmm(param.epsilon)])
-                vdw_terms.append(
-                    _VdwTerm(
-                        particle_index=atom_index,
-                        atom_type=atom_type,
-                        element=symbol,
-                        ff_row=param.ff_row,
-                    )
-                )
-            if use_harmonic:
-                exceptions_14 = _build_harmonic_exclusions(molecule, vdw_force)
-            else:
-                vdw_force.createExclusionsFromBonds([(bond.atom_i, bond.atom_j) for bond in molecule.bonds], 2)
-                _restrict_mm3_centers(vdw_force, vdw_assignments)
+        vdw_force, vdw_terms, exceptions_14 = _build_vdw_force(molecule, forcefield, use_harmonic=use_harmonic)
 
         # --- Assign CMAP correction terms (CHARMM backbone corrections) ---
         cmap_force, cmap_terms = _build_cmap_force(molecule, forcefield)
@@ -1060,35 +1076,15 @@ def _build_state(
                 "detected bonds, angles, torsions, vdW types, or CMAP pairs."
             )
 
-        if bond_terms:
-            system.addForce(bond_force)
-        else:
-            bond_force = None
-
-        if angle_terms:
-            system.addForce(angle_force)
-        else:
-            angle_force = None
-
-        if torsion_terms:
-            system.addForce(torsion_force)
-        else:
-            torsion_force = None
-
-        if vdw_terms:
-            system.addForce(vdw_force)
-        else:
-            vdw_force = None
-
-        if ub_terms:
-            system.addForce(ub_force)
-        else:
-            ub_force = None
-
-        if cmap_terms:
-            system.addForce(cmap_force)
-        else:
-            cmap_force = None
+        bond_force = bond_force if bond_terms else None
+        angle_force = angle_force if angle_terms else None
+        torsion_force = torsion_force if torsion_terms else None
+        vdw_force = vdw_force if vdw_terms else None
+        ub_force = ub_force if ub_terms else None
+        cmap_force = cmap_force if cmap_terms else None
+        for force in (bond_force, angle_force, torsion_force, vdw_force, ub_force, cmap_force):
+            if force is not None:
+                system.addForce(force)
 
         integrator, context = self._build_context(system, precision=precision)
         context.setPositions(self._positions(molecule))
@@ -1136,101 +1132,104 @@ def _update_params(self, state: _OpenMMState, forcefield: ForceField) -> None:
         use_harmonic = state.functional_form == FunctionalForm.HARMONIC
 
         if state.bond_force is not None:
-            for term, param in bond_assignments:
+            for bond_term, bond_param in bond_assignments:
                 if use_harmonic:
                     state.bond_force.setBondParameters(
-                        term.force_index,
-                        term.atom_i,
-                        term.atom_j,
-                        ang_to_nm(param.equilibrium),
-                        _bond_k_to_harmonic(param.force_constant),
+                        bond_term.force_index,
+                        bond_term.atom_i,
+                        bond_term.atom_j,
+                        ang_to_nm(bond_param.equilibrium),
+                        _bond_k_to_harmonic(bond_param.force_constant),
                     )
                 else:
                     state.bond_force.setBondParameters(
-                        term.force_index,
-                        term.atom_i,
-                        term.atom_j,
-                        [_bond_k_to_openmm(param.force_constant), ang_to_nm(param.equilibrium)],
+                        bond_term.force_index,
+                        bond_term.atom_i,
+                        bond_term.atom_j,
+                        [_bond_k_to_openmm(bond_param.force_constant), ang_to_nm(bond_param.equilibrium)],
                     )
             state.bond_force.updateParametersInContext(state.context)
 
         if state.angle_force is not None:
-            for term in state.angle_terms:
-                param = forcefield.match_angle(term.elements, env_id=term.env_id, ff_row=term.ff_row)
-                if param is None:
-                    raise ValueError(f"Updated force field is missing angle parameter for {term.elements}.")
+            for angle_term in state.angle_terms:
+                angle_param = forcefield.match_angle(
+                    angle_term.elements, env_id=angle_term.env_id, ff_row=angle_term.ff_row
+                )
+                if angle_param is None:
+                    raise ValueError(f"Updated force field is missing angle parameter for {angle_term.elements}.")
                 if use_harmonic:
                     state.angle_force.setAngleParameters(
-                        term.force_index,
-                        term.atom_i,
-                        term.atom_j,
-                        term.atom_k,
-                        np.deg2rad(float(param.equilibrium)),
-                        _angle_k_to_harmonic(param.force_constant),
+                        angle_term.force_index,
+                        angle_term.atom_i,
+                        angle_term.atom_j,
+                        angle_term.atom_k,
+                        np.deg2rad(float(angle_param.equilibrium)),
+                        _angle_k_to_harmonic(angle_param.force_constant),
                     )
                 else:
                     state.angle_force.setAngleParameters(
-                        term.force_index,
-                        term.atom_i,
-                        term.atom_j,
-                        term.atom_k,
-                        [_angle_k_to_openmm(param.force_constant), np.deg2rad(float(param.equilibrium))],
+                        angle_term.force_index,
+                        angle_term.atom_i,
+                        angle_term.atom_j,
+                        angle_term.atom_k,
+                        [_angle_k_to_openmm(angle_param.force_constant), np.deg2rad(float(angle_param.equilibrium))],
                     )
             state.angle_force.updateParametersInContext(state.context)
 
         if state.torsion_force is not None:
-            for term in state.torsion_terms:
+            for torsion_term in state.torsion_terms:
                 params = forcefield.match_torsion(
-                    term.elements, env_id=term.env_id, ff_row=term.ff_row, is_improper=term.is_improper
+                    torsion_term.elements,
+                    env_id=torsion_term.env_id,
+                    ff_row=torsion_term.ff_row,
+                    is_improper=torsion_term.is_improper,
                 )
-                matched = [p for p in params if p.periodicity == term.periodicity]
+                matched = [p for p in params if p.periodicity == torsion_term.periodicity]
                 if not matched:
                     raise ValueError(
                         f"Updated force field is missing torsion parameter for "
-                        f"{term.elements} periodicity={term.periodicity}."
+                        f"{torsion_term.elements} periodicity={torsion_term.periodicity}."
                     )
-                param = matched[0]
+                torsion_param = matched[0]
                 state.torsion_force.setTorsionParameters(
-                    term.force_index,
-                    term.atom_i,
-                    term.atom_j,
-                    term.atom_k,
-                    term.atom_l,
-                    param.periodicity,
-                    np.deg2rad(float(param.phase)),
-                    canonical_to_openmm_torsion_k(param.force_constant),
+                    torsion_term.force_index,
+                    torsion_term.atom_i,
+                    torsion_term.atom_j,
+                    torsion_term.atom_k,
+                    torsion_term.atom_l,
+                    torsion_param.periodicity,
+                    np.deg2rad(float(torsion_param.phase)),
+                    canonical_to_openmm_torsion_k(torsion_param.force_constant),
                 )
             state.torsion_force.updateParametersInContext(state.context)
 
         if state.vdw_force is not None:
             reinitialize = False
-            for term in state.vdw_terms:
-                param = forcefield.match_vdw(atom_type=term.atom_type, element=term.element, ff_row=term.ff_row)
-                if param is None:
+            for vdw_term in state.vdw_terms:
+                vdw_param = forcefield.match_vdw(
+                    atom_type=vdw_term.atom_type, element=vdw_term.element, ff_row=vdw_term.ff_row
+                )
+                if vdw_param is None:
                     raise ValueError(
-                        f"Updated force field is missing vdW parameter for {term.atom_type or term.element}."
+                        f"Updated force field is missing vdW parameter for {vdw_term.atom_type or vdw_term.element}."
                     )
                 if use_harmonic:
                     state.vdw_force.setParticleParameters(
-                        term.particle_index,
+                        vdw_term.particle_index,
                         0.0,
-                        _vdw_sigma_nm(param.radius),
-                        _vdw_epsilon_to_openmm(param.epsilon),
+                        _vdw_sigma_nm(vdw_param.radius),
+                        _vdw_epsilon_to_openmm(vdw_param.epsilon),
                     )
                 else:
                     state.vdw_force.setParticleParameters(
-                        term.particle_index,
-                        [_vdw_radius_to_openmm(param.radius), _vdw_epsilon_to_openmm(param.epsilon)],
+                        vdw_term.particle_index,
+                        [_vdw_radius_to_openmm(vdw_param.radius), _vdw_epsilon_to_openmm(vdw_param.epsilon)],
                     )
 
             # Recompute 1-4 exception params from updated particle params
             if use_harmonic and state.exceptions_14:
-                SCNB = 2.0
                 for exc in state.exceptions_14:
-                    _, sig1, eps1 = state.vdw_force.getParticleParameters(exc.particle_i)
-                    _, sig2, eps2 = state.vdw_force.getParticleParameters(exc.particle_j)
-                    sig_14 = 0.5 * (sig1 + sig2)
-                    eps_14 = (eps1 * eps2) ** 0.5 / SCNB
+                    sig_14, eps_14 = _exception_14_params(state.vdw_force, exc.particle_i, exc.particle_j)
                     old_epsilon = state.vdw_force.getExceptionParameters(exc.exception_index)[4]
                     reinitialize |= (old_epsilon == 0 * old_epsilon.unit) != (eps_14 == 0 * eps_14.unit)
                     state.vdw_force.setExceptionParameters(
@@ -1244,24 +1243,22 @@ def _update_params(self, state: _OpenMMState, forcefield: ForceField) -> None:
                 state.vdw_force.updateParametersInContext(state.context)
 
         if state.ub_force is not None:
-            for term in state.ub_terms:
-                param = forcefield.match_angle(term.elements, env_id=term.env_id, ff_row=term.ff_row)
-                if param is None:
-                    raise ValueError(f"Updated force field is missing UB parameter for angle {term.elements}.")
-                if param.ub_force_constant is None and param.ub_equilibrium is None:
-                    raise ValueError(f"Updated force field is missing UB parameter for angle {term.elements}.")
-                if param.ub_force_constant is None or param.ub_equilibrium is None:
+            for ub_term in state.ub_terms:
+                ub_param = forcefield.match_angle(ub_term.elements, env_id=ub_term.env_id, ff_row=ub_term.ff_row)
+                if ub_param is None or (ub_param.ub_force_constant is None and ub_param.ub_equilibrium is None):
+                    raise ValueError(f"Updated force field is missing UB parameter for angle {ub_term.elements}.")
+                if ub_param.ub_force_constant is None or ub_param.ub_equilibrium is None:
                     raise ValueError(
                         "Inconsistent Urey-Bradley parameters for angle "
-                        f"{term.elements} (env_id={term.env_id}, ff_row={term.ff_row}): "
+                        f"{ub_term.elements} (env_id={ub_term.env_id}, ff_row={ub_term.ff_row}): "
                         "both 'ub_force_constant' and 'ub_equilibrium' must be set or both must be None."
                     )
                 state.ub_force.setBondParameters(
-                    term.force_index,
-                    term.atom_i,
-                    term.atom_k,
-                    ang_to_nm(param.ub_equilibrium),
-                    _bond_k_to_harmonic(param.ub_force_constant),
+                    ub_term.force_index,
+                    ub_term.atom_i,
+                    ub_term.atom_k,
+                    ang_to_nm(ub_param.ub_equilibrium),
+                    _bond_k_to_harmonic(ub_param.ub_force_constant),
                 )
             state.ub_force.updateParametersInContext(state.context)
 
@@ -1312,8 +1309,8 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
 
         # --- Bonds: each bond param contributes (k, r0) ---
         bond_global_map: dict[int, tuple[str, str]] = {}
-        bond_k_factor = canonical_to_openmm_bond_k_nm(1.0) if use_harmonic else _bond_k_to_openmm(1.0)
-        for bp_idx, bp in enumerate(forcefield.bonds):
+        bond_k_factor = float(canonical_to_openmm_bond_k_nm(1.0) if use_harmonic else _bond_k_to_openmm(1.0))
+        for bp_idx, _bp in enumerate(forcefield.bonds):
             k_name = f"bond_k_{bp_idx}"
             r0_name = f"bond_r0_{bp_idx}"
             bond_global_map[bp_idx] = (k_name, r0_name)
@@ -1329,7 +1326,7 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
 
         for bp_idx, bp in enumerate(forcefield.bonds):
             k_name, r0_name = bond_global_map[bp_idx]
-            k_val = (
+            bond_k_val = float(
                 canonical_to_openmm_bond_k_nm(bp.force_constant)
                 if use_harmonic
                 else _bond_k_to_openmm(bp.force_constant)
@@ -1340,7 +1337,7 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
                 expr = f"{k_name}*dr10^2*(1-c3*dr10+c4*dr10^2);dr10=10*(r-{r0_name});c3={MM3_BOND_C3};c4={MM3_BOND_C4}"
             bf = mm.CustomBondForce(expr)
             bf.setForceGroup(0)
-            bf.addGlobalParameter(k_name, k_val)
+            bf.addGlobalParameter(k_name, bond_k_val)
             bf.addGlobalParameter(r0_name, ang_to_nm(bp.equilibrium))
             bf.addEnergyParameterDerivative(k_name)
             bf.addEnergyParameterDerivative(r0_name)
@@ -1355,8 +1352,8 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
 
         # --- Angles: each angle param contributes (k, theta0) ---
         angle_global_map: dict[int, tuple[str, str]] = {}
-        angle_k_factor = _angle_k_to_openmm(1.0)
-        for ap_idx, ap in enumerate(forcefield.angles):
+        angle_k_factor = float(_angle_k_to_openmm(1.0))
+        for ap_idx, _ap in enumerate(forcefield.angles):
             k_name = f"angle_k_{ap_idx}"
             t0_name = f"angle_t0_{ap_idx}"
             angle_global_map[ap_idx] = (k_name, t0_name)
@@ -1365,14 +1362,14 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
             grad_unit_factors.extend(
                 [
                     angle_k_factor,
-                    np.deg2rad(1.0),
+                    float(np.deg2rad(1.0)),
                 ]
             )
             pv_idx += 2
 
         for ap_idx, ap in enumerate(forcefield.angles):
             k_name, t0_name = angle_global_map[ap_idx]
-            k_val = _angle_k_to_openmm(ap.force_constant)
+            angle_k_val = float(_angle_k_to_openmm(ap.force_constant))
             if use_harmonic:
                 expr = f"{k_name}*(theta-{t0_name})^2"
             else:
@@ -1391,7 +1388,7 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
                 )
             af = mm.CustomAngleForce(expr)
             af.setForceGroup(1)
-            af.addGlobalParameter(k_name, k_val)
+            af.addGlobalParameter(k_name, angle_k_val)
             af.addGlobalParameter(t0_name, np.deg2rad(float(ap.equilibrium)))
             af.addEnergyParameterDerivative(k_name)
             af.addEnergyParameterDerivative(t0_name)
@@ -1408,8 +1405,8 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
 
         # --- Torsions: each torsion param (proper and improper) contributes (k,) ---
         torsion_global_map: dict[int, str] = {}
-        torsion_k_factor = canonical_to_openmm_torsion_k(1.0)
-        for tp_idx, tp in enumerate(forcefield.torsions):
+        torsion_k_factor = float(canonical_to_openmm_torsion_k(1.0))
+        for tp_idx, _tp in enumerate(forcefield.torsions):
             k_name = f"torsion_k_{tp_idx}"
             torsion_global_map[tp_idx] = k_name
             param_names.append(k_name)
@@ -1419,13 +1416,13 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
 
         for tp_idx, tp in enumerate(forcefield.torsions):
             k_name = torsion_global_map[tp_idx]
-            k_val = canonical_to_openmm_torsion_k(tp.force_constant)
+            torsion_k_val = float(canonical_to_openmm_torsion_k(tp.force_constant))
             phase_rad = np.deg2rad(float(tp.phase))
             n = tp.periodicity
             expr = f"{k_name}*(1+cos({n}*theta-{phase_rad:.15g}))"
             tf = mm.CustomTorsionForce(expr)
             tf.setForceGroup(1)
-            tf.addGlobalParameter(k_name, k_val)
+            tf.addGlobalParameter(k_name, torsion_k_val)
             tf.addEnergyParameterDerivative(k_name)
 
             for torsion in torsions_by_param.get(id(tp), []):
@@ -1444,41 +1441,7 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
         pv_idx += 2 * len(forcefield.vdws)
 
         if forcefield.vdws:
-            vdw_assignments = _collect_vdw_assignments(molecule, forcefield)
-            if use_harmonic:
-                vdw_force = mm.NonbondedForce()
-                vdw_force.setNonbondedMethod(mm.NonbondedForce.NoCutoff)
-                for _atom_index, _symbol, _atom_type, param in vdw_assignments:
-                    if param is None:
-                        vdw_force.addParticle(0.0, 0.0, 0.0)
-                    else:
-                        vdw_force.addParticle(
-                            0.0,
-                            _vdw_sigma_nm(param.radius),
-                            _vdw_epsilon_to_openmm(param.epsilon),
-                        )
-                _build_harmonic_exclusions(molecule, vdw_force)
-            else:
-                vdw_force = mm.CustomNonbondedForce(
-                    f"step(r - rc) * epsilon*(-{MM3_VDW_C}*(rv/r)^6 + {MM3_VDW_A}*exp(-{MM3_VDW_B}*r/rv))"
-                    f" + step(rc - r) * epsilon*{MM3_VDW_A}*exp(-{MM3_VDW_B}*rc/rv) * (rc/r)^12;"
-                    "rc=0.34*rv;"
-                    "rv=radius1+radius2;"
-                    "epsilon=sqrt(epsilon1*epsilon2)"
-                )
-                vdw_force.addPerParticleParameter("radius")
-                vdw_force.addPerParticleParameter("epsilon")
-                vdw_force.setNonbondedMethod(mm.CustomNonbondedForce.NoCutoff)
-                for _atom_index, _symbol, _atom_type, param in vdw_assignments:
-                    if param is None:
-                        vdw_force.addParticle([0.0, 0.0])
-                    else:
-                        vdw_force.addParticle(
-                            [_vdw_radius_to_openmm(param.radius), _vdw_epsilon_to_openmm(param.epsilon)]
-                        )
-                vdw_force.createExclusionsFromBonds([(b.atom_i, b.atom_j) for b in molecule.bonds], 2)
-                _restrict_mm3_centers(vdw_force, vdw_assignments)
-
+            vdw_force, _vdw_terms, _exceptions_14 = _build_vdw_force(molecule, forcefield, use_harmonic=use_harmonic)
             system.addForce(vdw_force)
 
         # --- Urey-Bradley: each UB angle contributes (ub_k, ub_r0) on the
@@ -1486,8 +1449,9 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
         # vector (after vdW), mirroring ParameterLayout ordering.  Modeled as a
         # harmonic bond so the energy and derivatives match energy()'s
         # HarmonicBondForce term. ---
-        ub_k_factor = canonical_to_openmm_bond_k_nm(1.0)
+        ub_k_factor = float(canonical_to_openmm_bond_k_nm(1.0))
         for ub_idx, ap in enumerate(forcefield._ub_angles):
+            assert ap.ub_force_constant is not None and ap.ub_equilibrium is not None
             k_name = f"ub_k_{ub_idx}"
             r0_name = f"ub_r0_{ub_idx}"
             param_names.extend([k_name, r0_name])
@@ -1502,7 +1466,7 @@ def _build_diff_state(self, state: _OpenMMState, forcefield: ForceField) -> _Ope
 
             ubf = mm.CustomBondForce(f"{k_name}*(r-{r0_name})^2")
             ubf.setForceGroup(0)
-            ubf.addGlobalParameter(k_name, canonical_to_openmm_bond_k_nm(ap.ub_force_constant))
+            ubf.addGlobalParameter(k_name, float(canonical_to_openmm_bond_k_nm(ap.ub_force_constant)))
             ubf.addGlobalParameter(r0_name, ang_to_nm(ap.ub_equilibrium))
             ubf.addEnergyParameterDerivative(k_name)
             ubf.addEnergyParameterDerivative(r0_name)
@@ -1739,10 +1703,6 @@ def __init__(
         self._backend = backend
         self._state = state
 
-    def _ff_for(self, parameters: np.ndarray) -> ForceField:
-        vec = self._validate_vector(parameters)
-        return self.layout.replace(self.force_field, vec)
-
     def _openmm_system(self) -> object:
         """Return this session's OpenMM ``System`` (I/O-boundary accessor).
 
@@ -1751,15 +1711,16 @@ def _openmm_system(self) -> object:
         """
         return self._state.system
 
-    def _energy(self, request: EnergyRequest) -> EnergyResult:  # type: ignore[override]
+    def _energy(self, request: EnergyRequest | ReferenceEnergyRequest) -> EnergyResult:
+        assert isinstance(request, EnergyRequest)
         ff = self._ff_for(request.parameters)
         try:
             value = self._backend._evaluate_energy(self._state, ff)
         except Exception as exc:  # noqa: BLE001
             raise EvaluationError(f"OpenMM energy evaluation failed: {exc}") from exc
-        return EnergyResult(energy=float(value), unit=EnergyUnit.KCAL_PER_MOL, provenance=self._info.provenance)
+        return EnergyResult(energy=float(value), unit=EnergyUnit.KCAL_PER_MOL, provenance=self._provenance)
 
-    def _minimize(self, request: MinimizationRequest) -> GeometryResult:  # type: ignore[override]
+    def _minimize(self, request: MinimizationRequest) -> GeometryResult:
         ff = self._ff_for(request.parameters)
         tolerance = request.tolerance if request.tolerance is not None else 1.0
         max_iterations = request.max_iterations if request.max_iterations is not None else 200
@@ -1778,35 +1739,29 @@ def _minimize(self, request: MinimizationRequest) -> GeometryResult:  # type: ig
             symbols=tuple(atoms),
             coordinates=readonly_array(coords),
             coordinate_unit=LengthUnit.ANGSTROM,
-            provenance=self._info.provenance,
+            provenance=self._provenance,
         )
 
-    def _hessian(self, request: HessianRequest) -> HessianResult:  # type: ignore[override]
+    def _raw_hessian(self, ff: ForceField) -> np.ndarray:
+        self._backend._update_params(self._state, ff)
+        return self._backend._evaluate_hessian(self._state)
+
+    def _hessian(self, request: HessianRequest | ReferenceHessianRequest) -> HessianResult:
+        assert isinstance(request, HessianRequest)
         ff = self._ff_for(request.parameters)
         try:
-            self._backend._update_params(self._state, ff)
-            hess = self._backend._evaluate_hessian(self._state)
+            hess = self._raw_hessian(ff)
         except Exception as exc:  # noqa: BLE001
             raise EvaluationError(f"OpenMM Hessian evaluation failed: {exc}") from exc
         return HessianResult(
-            hessian=readonly_array(hess), unit=HessianUnit.HARTREE_PER_BOHR2, provenance=self._info.provenance
+            hessian=readonly_array(hess), unit=HessianUnit.HARTREE_PER_BOHR2, provenance=self._provenance
         )
 
-    def _frequencies(self, request: FrequencyRequest) -> FrequencyResult:  # type: ignore[override]
-        from q2mm.models.hessian import hessian_to_frequencies
-
-        ff = self._ff_for(request.parameters)
-        try:
-            self._backend._update_params(self._state, ff)
-            hess_au = self._backend._evaluate_hessian(self._state)
-            freqs = hessian_to_frequencies(hess_au, list(self.molecule.symbols), on_error=request.on_error)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"OpenMM frequency evaluation failed: {exc}") from exc
-        return FrequencyResult(
-            frequencies=readonly_array(freqs), unit=FrequencyUnit.INVERSE_CM, provenance=self._info.provenance
-        )
+    def _frequencies(self, request: FrequencyRequest | ReferenceFrequencyRequest) -> FrequencyResult:
+        assert isinstance(request, FrequencyRequest)
+        return self._mm_frequencies(request, backend_label="OpenMM")
 
-    def _parameter_gradient(self, request: ParameterGradientRequest) -> ParameterGradientResult:  # type: ignore[override]
+    def _parameter_gradient(self, request: ParameterGradientRequest) -> ParameterGradientResult:
         ff = self._ff_for(request.parameters)
         try:
             energy, grad = self._backend._evaluate_param_grad(self._state, ff)
@@ -1816,5 +1771,5 @@ def _parameter_gradient(self, request: ParameterGradientRequest) -> ParameterGra
             energy=float(energy),
             gradient=readonly_array(grad),
             unit=EnergyUnit.KCAL_PER_MOL,
-            provenance=self._info.provenance,
+            provenance=self._provenance,
         )
diff --git a/q2mm/backends/mm/tinker.py b/q2mm/backends/mm/tinker.py
index ea9ebc33..81789699 100644
--- a/q2mm/backends/mm/tinker.py
+++ b/q2mm/backends/mm/tinker.py
@@ -35,7 +35,6 @@
     EvaluationError,
     FrequencyRequest,
     FrequencyResult,
-    FrequencyUnit,
     GeometryResult,
     HessianRequest,
     HessianResult,
@@ -44,8 +43,12 @@
     MinimizationRequest,
     PreparationError,
     PreparationRequest,
+    ReferenceEnergyRequest,
+    ReferenceFrequencyRequest,
+    ReferenceHessianRequest,
     readonly_array,
 )
+from q2mm.backends.mm._term_support import _require_mm_force_field
 from q2mm.constants import (
     DEFAULT_BOND_TOLERANCE,
     MM3_BOND_C3,
@@ -203,12 +206,13 @@ def __init__(
         """
         from q2mm.backends.contracts import BackendUnavailableError
 
-        self._tinker_dir = tinker_dir or _find_tinker_dir()
+        resolved_dir = tinker_dir or _find_tinker_dir()
         self._bond_tolerance = bond_tolerance
-        if self._tinker_dir is None:
+        if resolved_dir is None:
             raise BackendUnavailableError(
                 "Tinker not found. Install from https://dasher.wustl.edu/tinker/ or pass tinker_dir parameter."
             )
+        self._tinker_dir: str = resolved_dir
 
         if params_file is None:
             # Try common locations for MM3 params
@@ -220,18 +224,30 @@ def __init__(
                 if os.path.isfile(c):
                     params_file = c
                     break
-        self._params_file = params_file
-        if self._params_file is None:
+        if params_file is None:
             raise BackendUnavailableError(
                 "MM3 parameter file not found. Provide params_file parameter "
                 "or place mm3.prm alongside the Tinker bin directory."
             )
+        self._params_file: str = params_file
 
     @property
     def info(self) -> BackendInfo:
         """Immutable capability declaration for this backend."""
         return _TINKER_INFO
 
+    def _uses_template_export(self, forcefield: ForceField) -> bool:
+        """Validate MM3 form and Tinker export terms, returning whether template export applies."""
+        _validate_form(forcefield, _TINKER_INFO)
+        use_template = forcefield.source_format == "tinker_prm" and bool(forcefield.source_path or self._params_file)
+        _validate_tinker_export_terms(
+            forcefield,
+            supports_reduction=use_template,
+            supports_placeholder_elements=use_template,
+            error_type=PreparationError,
+        )
+        return use_template
+
     def prepare(self, request: PreparationRequest) -> PreparedTinker:
         """Build a prepared session for one training case.
 
@@ -247,31 +263,21 @@ def prepare(self, request: PreparationRequest) -> PreparedTinker:
                 form or populated terms are unsupported.
 
         """
-        if request.force_field is None:
-            raise PreparationError("Tinker requires a base ForceField in the PreparationRequest.")
-        _validate_form(request.force_field, _TINKER_INFO)
-        use_template = request.force_field.source_format == "tinker_prm" and bool(
-            request.force_field.source_path or self._params_file
-        )
-        _validate_tinker_export_terms(
-            request.force_field,
-            supports_reduction=use_template,
-            supports_placeholder_elements=use_template,
-            error_type=PreparationError,
-        )
+        force_field = _require_mm_force_field(request, _TINKER_INFO, backend="Tinker")
+        use_template = self._uses_template_export(force_field)
         try:
             atom_types = self._molecule_type_numbers(request.molecule)
             self._tinker_xyz_lines(request.molecule, atom_types)
             if not use_template:
-                self._standalone_prm_lines(request.force_field, list(request.molecule.symbols), atom_types)
+                self._standalone_prm_lines(force_field, list(request.molecule.symbols), atom_types)
         except ValueError as exc:
             raise PreparationError(f"Tinker preparation failed: {exc}") from exc
-        layout = ParameterLayout.from_force_field(request.force_field)
+        layout = ParameterLayout.from_force_field(force_field)
         return PreparedTinker(
             backend=self,
             case_id=request.case_id,
             molecule=request.molecule,
-            force_field=request.force_field,
+            force_field=force_field,
             layout=layout,
         )
 
@@ -289,14 +295,7 @@ def _write_tinker_xyz(self, molecule: Molecule, forcefield: ForceField, workdir:
             str: Path to the Tinker XYZ file.
 
         """
-        _validate_form(forcefield, _TINKER_INFO)
-        use_template = forcefield.source_format == "tinker_prm" and bool(forcefield.source_path or self._params_file)
-        _validate_tinker_export_terms(
-            forcefield,
-            supports_reduction=use_template,
-            supports_placeholder_elements=use_template,
-            error_type=PreparationError,
-        )
+        use_template = self._uses_template_export(forcefield)
 
         atoms = list(molecule.symbols)
         atom_type_numbers = self._molecule_type_numbers(molecule)
@@ -336,8 +335,9 @@ def _write_tinker_xyz(self, molecule: Molecule, forcefield: ForceField, workdir:
     def _molecule_type_numbers(molecule: Molecule) -> list[int]:
         """Resolve actual XYZ types without guessing an unknown element's class."""
         default_types = {"C": 1, "H": 5, "F": 11, "Cl": 12, "Br": 13, "N": 8, "O": 6, "S": 15, "P": 25}
+        atom_types = molecule.atom_types if molecule.atom_types is not None else molecule.symbols
         numbers = []
-        for atom, atom_type in zip(molecule.symbols, molecule.atom_types, strict=True):
+        for atom, atom_type in zip(molecule.symbols, atom_types, strict=True):
             number = _tinker_atom_type_number(atom_type)
             if number is None:
                 if atom not in default_types:
@@ -351,7 +351,7 @@ def _molecule_type_numbers(molecule: Molecule) -> list[int]:
     def _tinker_xyz_lines(molecule: Molecule, atom_type_numbers: list[int]) -> list[str]:
         """Stage XYZ records so invalid coordinates cannot replace any inputs."""
         bonds: dict[int, list[int]] = {i: [] for i in range(molecule.n_atoms)}
-        for bond in molecule.bonds:
+        for bond in molecule.bonds or ():
             bonds[bond.atom_i].append(bond.atom_j)
             bonds[bond.atom_j].append(bond.atom_i)
         lines = [f"     {molecule.n_atoms}  Q2MM Tinker input\n"]
@@ -485,18 +485,18 @@ def _check_elements(elements: tuple[str, ...], label: str) -> None:
             _require_unique_tinker_types(identities, "bond", types)
             t1, t2 = types
             _require_finite_tinker_values(bond.force_constant, bond.equilibrium)
-            native_k = canonical_to_mm3_bond_k(bond.force_constant)
-            _require_finite_tinker_values(native_k)
-            lines.append(f"bond   {t1:5d} {t2:5d}         {native_k:8.4f}   {bond.equilibrium:8.4f}\n")
+            bond_k = canonical_to_mm3_bond_k(bond.force_constant)
+            _require_finite_tinker_values(bond_k)
+            lines.append(f"bond   {t1:5d} {t2:5d}         {bond_k:8.4f}   {bond.equilibrium:8.4f}\n")
 
         for angle in ff.angles:
             types = tuple(elem_to_type[e] for e in angle.elements)
             _require_unique_tinker_types(identities, "angle", types)
             t1, t2, t3 = types
             _require_finite_tinker_values(angle.force_constant, angle.equilibrium)
-            native_k = canonical_to_mm3_angle_k(angle.force_constant)
-            _require_finite_tinker_values(native_k)
-            lines.append(f"angle  {t1:5d} {t2:5d} {t3:5d}         {native_k:8.4f}   {angle.equilibrium:8.4f}\n")
+            angle_k = canonical_to_mm3_angle_k(angle.force_constant)
+            _require_finite_tinker_values(angle_k)
+            lines.append(f"angle  {t1:5d} {t2:5d} {t3:5d}         {angle_k:8.4f}   {angle.equilibrium:8.4f}\n")
 
         torsion_groups: dict[tuple[int, ...], list[TorsionParam]] = {}
         for tp in ff.torsions:
@@ -507,9 +507,9 @@ def _check_elements(elements: tuple[str, ...], label: str) -> None:
             for tp in tps:
                 _require_finite_tinker_values(tp.force_constant, tp.phase)
                 raw_parts.extend((str(tp.force_constant), str(tp.phase), str(tp.periodicity)))
-            terms = _tinker_torsion_terms(raw_parts, len(lines) + 1)
+            torsion_terms = _tinker_torsion_terms(raw_parts, len(lines) + 1)
             parts = []
-            for amplitude, phase, fold in sorted(terms, key=lambda term: term[2]):
+            for amplitude, phase, fold in sorted(torsion_terms, key=lambda term: term[2]):
                 parts.extend([f"{amplitude:8.4f}", f"{phase:6.1f}", f" {fold}"])
             lines.append(f"torsion {types[0]:4d} {types[1]:4d} {types[2]:4d} {types[3]:4d}  {'  '.join(parts)}\n")
 
@@ -729,27 +729,6 @@ def _evaluate_hessian(self, structure: Molecule, forcefield: ForceField) -> np.n
             # Tinker outputs Hessian in kcal/(mol·Å²); convert to Hartree/Bohr²
             return hessian * KCALMOLA2_TO_HESSIAN_AU
 
-    def _evaluate_frequencies(
-        self, structure: Molecule, forcefield: ForceField, on_error: str = "raise"
-    ) -> list[float]:
-        """Calculate vibrational frequencies in cm⁻¹.
-
-        Args:
-            structure (Molecule): Molecule to evaluate.
-            forcefield (ForceField): Force field with the parameter values.
-            on_error: Forwarded to
-                :func:`~q2mm.models.hessian.hessian_to_frequencies`.
-
-        Returns:
-            list[float]: Vibrational frequencies in cm⁻¹.
-
-        """
-        from q2mm.models.hessian import hessian_to_frequencies
-
-        hessian_au = self._evaluate_hessian(structure, forcefield)
-        symbols = list(structure.symbols)
-        return hessian_to_frequencies(hessian_au, symbols, on_error=on_error)
-
 
 class PreparedTinker(AbstractPreparedBackend):
     """Prepared Tinker session for a single training case.
@@ -777,16 +756,13 @@ def __init__(
         )
         self._backend = backend
 
-    def _ff_for(self, parameters: np.ndarray) -> ForceField:
-        vec = self._validate_vector(parameters)
-        return self.layout.replace(self.force_field, vec)
-
-    def _energy(self, request: EnergyRequest) -> EnergyResult:  # type: ignore[override]
+    def _energy(self, request: EnergyRequest | ReferenceEnergyRequest) -> EnergyResult:
+        assert isinstance(request, EnergyRequest)
         ff = self._ff_for(request.parameters)
         value = self._backend._evaluate_energy(self.molecule, ff)
         return EnergyResult(energy=float(value), unit=EnergyUnit.KCAL_PER_MOL, provenance=_TINKER_PROVENANCE)
 
-    def _minimize(self, request: MinimizationRequest) -> GeometryResult:  # type: ignore[override]
+    def _minimize(self, request: MinimizationRequest) -> GeometryResult:
         ff = self._ff_for(request.parameters)
         rms_grad = request.tolerance if request.tolerance is not None else 0.01
         energy, atoms, coords = self._backend._evaluate_minimize(
@@ -801,19 +777,17 @@ def _minimize(self, request: MinimizationRequest) -> GeometryResult:  # type: ig
             provenance=_TINKER_PROVENANCE,
         )
 
-    def _hessian(self, request: HessianRequest) -> HessianResult:  # type: ignore[override]
+    def _raw_hessian(self, ff: ForceField) -> np.ndarray:
+        return self._backend._evaluate_hessian(self.molecule, ff)
+
+    def _hessian(self, request: HessianRequest | ReferenceHessianRequest) -> HessianResult:
+        assert isinstance(request, HessianRequest)
         ff = self._ff_for(request.parameters)
-        hess = self._backend._evaluate_hessian(self.molecule, ff)
+        hess = self._raw_hessian(ff)
         return HessianResult(
             hessian=readonly_array(hess), unit=HessianUnit.HARTREE_PER_BOHR2, provenance=_TINKER_PROVENANCE
         )
 
-    def _frequencies(self, request: FrequencyRequest) -> FrequencyResult:  # type: ignore[override]
-        ff = self._ff_for(request.parameters)
-        try:
-            freqs = self._backend._evaluate_frequencies(self.molecule, ff, on_error=request.on_error)
-        except Exception as exc:  # noqa: BLE001
-            raise EvaluationError(f"Tinker frequency evaluation failed: {exc}") from exc
-        return FrequencyResult(
-            frequencies=readonly_array(freqs), unit=FrequencyUnit.INVERSE_CM, provenance=_TINKER_PROVENANCE
-        )
+    def _frequencies(self, request: FrequencyRequest | ReferenceFrequencyRequest) -> FrequencyResult:
+        assert isinstance(request, FrequencyRequest)
+        return self._mm_frequencies(request, backend_label="Tinker")
diff --git a/q2mm/backends/reference/__init__.py b/q2mm/backends/reference/__init__.py
index b708293b..7efb0305 100644
--- a/q2mm/backends/reference/__init__.py
+++ b/q2mm/backends/reference/__init__.py
@@ -1 +1,17 @@
 """Reference-data backend implementations."""
+
+from __future__ import annotations
+
+from q2mm.backends.contracts import PreparationError, PreparationRequest
+from q2mm.models.molecule import Molecule
+
+
+def _validate_reference_prepare(request: PreparationRequest, *, backend: str) -> Molecule:
+    """Validate that a reference preparation request carries only a Molecule."""
+    if request.force_field is not None:
+        raise PreparationError(f"{backend} reference preparation does not accept a force_field.")
+    if request.options:
+        raise PreparationError(f"{backend} reference preparation does not accept per-case options.")
+    if not isinstance(request.molecule, Molecule):
+        raise PreparationError(f"{backend} reference preparation requires a q2mm Molecule.")
+    return request.molecule
diff --git a/q2mm/backends/reference/ase.py b/q2mm/backends/reference/ase.py
index 317f0506..354889c7 100644
--- a/q2mm/backends/reference/ase.py
+++ b/q2mm/backends/reference/ase.py
@@ -36,6 +36,7 @@
     ReferenceEnergyRequest,
     readonly_array,
 )
+from q2mm.backends.reference import _validate_reference_prepare
 from q2mm.models.molecule import Molecule
 
 _ase: Any = None
@@ -44,17 +45,13 @@
 _HARTREE_EV: Any = None
 
 try:
-    import ase as _ase_import
-    from ase import Atoms as _Atoms_import
-    from ase.units import Bohr as _BOHR_ANGSTROM_IMPORT
-    from ase.units import Hartree as _HARTREE_EV_IMPORT
+    import ase as _ase_mod
+    from ase import Atoms as _Atoms_cls
+    from ase.units import Bohr as _BOHR_VAL, Hartree as _HARTREE_VAL
+
+    _ase, _Atoms, _BOHR_ANGSTROM, _HARTREE_EV = _ase_mod, _Atoms_cls, _BOHR_VAL, _HARTREE_VAL
 except ImportError:
     pass
-else:
-    _ase = _ase_import
-    _Atoms = _Atoms_import
-    _BOHR_ANGSTROM = _BOHR_ANGSTROM_IMPORT
-    _HARTREE_EV = _HARTREE_EV_IMPORT
 
 _CONVERSION_VERSION = 1
 
@@ -129,19 +126,14 @@ def _runtime_capabilities(calculator: object) -> tuple[frozenset[Capability], tu
     return frozenset(capabilities), implemented
 
 
-def _periodic_state_reason(atoms: object) -> str:
-    """Return a reason when an Atoms-like object carries PBC or a nonzero cell."""
+def _unsupported_atoms_state_reason(atoms: object) -> str:
+    """Return a reason when an Atoms-like object carries PBC, a cell, or per-atom charge/spin."""
     pbc = np.asarray(getattr(atoms, "pbc"), dtype=bool)
     cell = np.asarray(getattr(atoms, "cell"), dtype=float)
     if np.any(pbc):
         return "periodic boundary conditions are enabled"
     if cell.size and np.any(cell != 0.0):
         return "a nonzero simulation cell is present"
-    return ""
-
-
-def _per_atom_state_reason(atoms: object) -> str:
-    """Return a reason when unsupported charge or spin arrays are present."""
     for name in ("initial_charges", "initial_magmoms"):
         if getattr(atoms, "has")(name):
             return f"unsupported per-atom {name} state is present"
@@ -151,7 +143,7 @@ def _per_atom_state_reason(atoms: object) -> str:
 def _validate_calculator_owned_atoms(calculator: object) -> None:
     try:
         atoms = getattr(calculator, "atoms", None)
-        reason = "" if atoms is None else _periodic_state_reason(atoms) or _per_atom_state_reason(atoms)
+        reason = "" if atoms is None else _unsupported_atoms_state_reason(atoms)
     except Exception as exc:
         raise BackendConfigurationError("ASE calculator-owned Atoms state could not be validated.") from exc
     if reason:
@@ -220,12 +212,7 @@ def info(self) -> BackendInfo:
 
     def prepare(self, request: PreparationRequest) -> PreparedASE:
         """Build an isolated non-periodic ASE Atoms/calculator session."""
-        if request.force_field is not None:
-            raise PreparationError("ASE reference preparation does not accept a force_field.")
-        if request.options:
-            raise PreparationError("ASE reference preparation does not accept per-case options.")
-        if not isinstance(request.molecule, Molecule):
-            raise PreparationError("ASE reference preparation requires a q2mm Molecule.")
+        molecule = _validate_reference_prepare(request, backend="ASE")
 
         calculator = _deepcopy_calculator(self._calculator_template, preparing=True)
         try:
@@ -235,7 +222,6 @@ def prepare(self, request: PreparationRequest) -> PreparedASE:
         if capabilities != self._capabilities or properties != self._implemented_properties:
             raise PreparationError("Copying the ASE calculator changed its runtime capabilities.")
 
-        molecule = request.molecule
         try:
             atoms = _Atoms(
                 symbols=list(molecule.symbols),
@@ -247,22 +233,16 @@ def prepare(self, request: PreparationRequest) -> PreparedASE:
                     "q2mm_multiplicity": molecule.multiplicity,
                 },
             )
-            reason = _periodic_state_reason(atoms)
-            if reason:
-                raise ValueError(reason)
-            reason = _per_atom_state_reason(atoms)
+            reason = _unsupported_atoms_state_reason(atoms)
             if reason:
                 raise ValueError(reason)
             atoms.calc = calculator
-            reason = _periodic_state_reason(atoms)
-            if reason:
-                raise ValueError(f"calculator attachment introduced {reason}")
-            reason = _per_atom_state_reason(atoms)
+            reason = _unsupported_atoms_state_reason(atoms)
             if reason:
                 raise ValueError(f"calculator attachment introduced {reason}")
             owned_atoms = getattr(calculator, "atoms", None)
             if owned_atoms is not None:
-                reason = _periodic_state_reason(owned_atoms) or _per_atom_state_reason(owned_atoms)
+                reason = _unsupported_atoms_state_reason(owned_atoms)
                 if reason:
                     raise ValueError(f"calculator attachment introduced {reason} in calculator-owned state")
         except Exception as exc:
@@ -363,15 +343,12 @@ def _validate_native_state(self) -> None:
                 raise ValueError("q2mm charge metadata changed")
             if self._atoms.info.get("q2mm_multiplicity") != self.molecule.multiplicity:
                 raise ValueError("q2mm multiplicity metadata changed")
-            reason = _periodic_state_reason(self._atoms)
-            if reason:
-                raise ValueError(reason)
-            reason = _per_atom_state_reason(self._atoms)
+            reason = _unsupported_atoms_state_reason(self._atoms)
             if reason:
                 raise ValueError(reason)
             owned_atoms = getattr(self._calculator, "atoms", None)
             if owned_atoms is not None:
-                reason = _periodic_state_reason(owned_atoms) or _per_atom_state_reason(owned_atoms)
+                reason = _unsupported_atoms_state_reason(owned_atoms)
                 if reason:
                     raise ValueError(f"calculator-owned Atoms state has {reason}")
         except Exception as exc:
diff --git a/q2mm/backends/reference/psi4.py b/q2mm/backends/reference/psi4.py
index e265145e..a11efeb5 100644
--- a/q2mm/backends/reference/psi4.py
+++ b/q2mm/backends/reference/psi4.py
@@ -175,15 +175,16 @@ def _evaluate_energy(self, molecule: Molecule) -> float:
         mol = self._load_molecule(molecule)
         return float(_psi4.energy(self._method, molecule=mol))
 
-    def _evaluate_hessian(self, molecule: Molecule) -> np.ndarray:
+    def _frequency_wfn(self, molecule: Molecule) -> Any:
         mol = self._load_molecule(molecule)
         _, wfn = _psi4.frequency(self._method, molecule=mol, return_wfn=True)
-        return np.array(wfn.hessian())
+        return wfn
+
+    def _evaluate_hessian(self, molecule: Molecule) -> np.ndarray:
+        return np.array(self._frequency_wfn(molecule).hessian())
 
     def _evaluate_frequencies(self, molecule: Molecule) -> list[float]:
-        mol = self._load_molecule(molecule)
-        _, wfn = _psi4.frequency(self._method, molecule=mol, return_wfn=True)
-        return list(np.array(wfn.frequencies()))
+        return list(np.array(self._frequency_wfn(molecule).frequencies()))
 
     def _evaluate_optimize(self, molecule: Molecule, opt_type: str) -> tuple[float, list[str], np.ndarray]:
         mol = self._load_molecule(molecule)
diff --git a/q2mm/backends/reference/qcengine.py b/q2mm/backends/reference/qcengine.py
index 4f966aad..851492a9 100644
--- a/q2mm/backends/reference/qcengine.py
+++ b/q2mm/backends/reference/qcengine.py
@@ -31,13 +31,13 @@
     HessianRequest,
     HessianResult,
     HessianUnit,
-    PreparationError,
     PreparationRequest,
     ReferenceCoordinateGradientRequest,
     ReferenceEnergyRequest,
     ReferenceHessianRequest,
     readonly_array,
 )
+from q2mm.backends.reference import _validate_reference_prepare
 from q2mm.constants import BOHR_TO_ANG
 from q2mm.models.molecule import Molecule
 
@@ -310,13 +310,8 @@ def info(self) -> BackendInfo:
 
     def prepare(self, request: PreparationRequest) -> PreparedQCEngine:
         """Create a reusable reference session without executing a job."""
-        if request.force_field is not None:
-            raise PreparationError("QCEngine reference preparation does not accept a force_field.")
-        if request.options:
-            raise PreparationError("QCEngine reference preparation does not accept per-case options.")
-        if not isinstance(request.molecule, Molecule):
-            raise PreparationError("QCEngine reference preparation requires a q2mm Molecule.")
-        return PreparedQCEngine(backend=self, case_id=request.case_id, molecule=request.molecule)
+        molecule = _validate_reference_prepare(request, backend="QCEngine")
+        return PreparedQCEngine(backend=self, case_id=request.case_id, molecule=molecule)
 
     def _build_specification(self, driver: str) -> Any:
         model: dict[str, object] = {"method": self._method}
diff --git a/test/test_backend_contracts.py b/test/test_backend_contracts.py
index 54432fef..994a0060 100644
--- a/test/test_backend_contracts.py
+++ b/test/test_backend_contracts.py
@@ -841,6 +841,122 @@ def test_coordinate_gradient_and_hessian_provenance_contracts() -> None:
         molecule_provenance.source_details["new"] = "value"  # type: ignore[index]
 
 
+def test_mm_frequencies_penalty_mode_recovers_from_nonfinite_raw_hessian() -> None:
+    """FrequencyRequest(on_error='penalty') converts non-finite raw Hessians to PENALTY_FREQUENCY."""
+    from q2mm.models.hessian import PENALTY_FREQUENCY
+
+    mol, ff, layout = _methane()
+    prov = BackendProvenance(backend="nan-hess", role=BackendRole.MM)
+    info = BackendInfo(
+        name="nan-hess",
+        role=BackendRole.MM,
+        capabilities=frozenset({Capability.HESSIAN, Capability.FREQUENCIES}),
+        functional_forms=frozenset({"harmonic"}),
+        provenance=prov,
+    )
+
+    class _NanHessianPrepared(AbstractPreparedBackend):
+        def _raw_hessian(self, _ff: ForceField) -> np.ndarray:
+            n3 = 3 * len(self.molecule.symbols)
+            hess = np.eye(n3)
+            hess[0, 0] = np.nan
+            return hess
+
+        def _hessian(self, request: HessianRequest | ReferenceHessianRequest) -> HessianResult:
+            assert isinstance(request, HessianRequest)
+            hess = self._raw_hessian(self._ff_for(request.parameters))
+            return HessianResult(
+                hessian=hess,
+                unit=HessianUnit.HARTREE_PER_BOHR2,
+                provenance=self._provenance,
+            )
+
+        def _frequencies(self, request: FrequencyRequest | object) -> FrequencyResult:
+            assert isinstance(request, FrequencyRequest)
+            return self._mm_frequencies(request, backend_label="nan-hess")
+
+    prepared = _NanHessianPrepared(info=info, case_id="0", molecule=mol, force_field=ff, layout=layout)
+    vec = layout.vector(ff)
+
+    # Strict HessianResult validation still rejects non-finite Hessians.
+    with pytest.raises(EvaluationError, match="non-finite"):
+        prepared.hessian(HessianRequest(parameters=vec))
+
+    # Default on_error="raise" still raises EvaluationError on non-finite Hessians.
+    with pytest.raises(EvaluationError, match="frequency evaluation failed"):
+        prepared.frequencies(FrequencyRequest(parameters=vec, on_error="raise"))
+
+    # Penalty mode returns PENALTY_FREQUENCY rather than failing in HessianResult.
+    freq_res = prepared.frequencies(FrequencyRequest(parameters=vec, on_error="penalty"))
+    assert freq_res.unit is FrequencyUnit.INVERSE_CM
+    assert np.all(freq_res.frequencies == PENALTY_FREQUENCY)
+    assert freq_res.frequencies.shape == (3 * len(mol.symbols),)
+
+
+def test_shared_mm_helpers_openmm_tinker_and_jax_topology() -> None:
+    """Directly verify shared OpenMM 1-4 scaling, Tinker template selection, and JAX topology matching."""
+    from dataclasses import replace
+    from pathlib import Path
+
+    from q2mm.backends.contracts import PreparationError
+    from q2mm.backends.mm._jax_common import match_mm_topology
+    from q2mm.backends.mm.openmm import _exception_14_params
+    from q2mm.backends.mm.tinker import TinkerBackend
+    from q2mm.models.forcefield import AngleParam, BondParam, TorsionParam, VdwParam
+
+    class _FakeVdwForce:
+        def getParticleParameters(self, idx: int) -> tuple[float, float, float]:
+            return (0.0, 0.30, 0.16) if idx == 0 else (0.0, 0.50, 0.04)
+
+    sig_14, eps_14 = _exception_14_params(_FakeVdwForce(), 0, 1)
+    assert sig_14 == pytest.approx(0.40)
+    assert eps_14 == pytest.approx(((0.16 * 0.04) ** 0.5) / 2.0)
+
+    tinker_backend = object.__new__(TinkerBackend)
+    tinker_backend._params_file = "default_mm3.prm"
+    standalone_ff = ForceField(
+        functional_form=FunctionalForm.MM3,
+        bonds=(BondParam(("C", "H"), 1.1, 100.0),),
+        angles=(AngleParam(("H", "C", "H"), 109.5, 50.0),),
+        vdws=(VdwParam("H", 1.2, 0.1),),
+    )
+    assert tinker_backend._uses_template_export(standalone_ff) is False
+    template_ff = replace(
+        standalone_ff,
+        source_format="tinker_prm",
+        source_path=Path("template.prm"),
+        vdws=(VdwParam("H", 1.2, 0.1, reduction=0.92),),
+    )
+    assert tinker_backend._uses_template_export(template_ff) is True
+    with pytest.raises(PreparationError, match="reduction"):
+        tinker_backend._uses_template_export(replace(template_ff, source_format=""))
+    with pytest.raises(PreparationError, match="functional form"):
+        tinker_backend._uses_template_export(replace(standalone_ff, functional_form=FunctionalForm.HARMONIC))
+
+    mol = Molecule(
+        symbols=("B", "H", "H", "H"),
+        geometry=np.array([[0.0, 0.0, 0.0], [1.2, 0.0, 0.0], [-0.6, 1.04, 0.0], [-0.6, -1.04, 0.0]]),
+        atom_types=("B", "H", "H", "H"),
+    )
+    ff_with_improper = ForceField(
+        functional_form=FunctionalForm.HARMONIC,
+        bonds=(BondParam(("B", "H"), 1.2, 100.0),),
+        angles=(AngleParam(("H", "B", "H"), 120.0, 50.0),),
+        torsions=(TorsionParam(("H", "H", "B", "H"), force_constant=1.0, is_improper=True),),
+        vdws=(VdwParam("B", 1.5, 0.05), VdwParam("H", 1.2, 0.02)),
+        nonbonded_excluded_atom_types=frozenset({"B"}),
+    )
+    topo_jax = match_mm_topology(mol, ff_with_improper, include_impropers=True, mask_excluded_vdw=False)
+    topo_jaxmd = match_mm_topology(mol, ff_with_improper, include_impropers=False, mask_excluded_vdw=True)
+    assert len(topo_jax.matched_bonds) == 3
+    assert len(topo_jax.matched_angles) == 3
+    assert len(topo_jax.torsion_indices) > 0
+    assert len(topo_jaxmd.torsion_indices) == 0
+    assert topo_jax.explicitly_excluded_atoms == frozenset({0})
+    assert topo_jax.atom_vdw_map[0] >= 0
+    assert topo_jaxmd.atom_vdw_map[0] == -1
+
+
 # ---------------------------------------------------------------------------
 # Preparation reuse: JAX objective executor prepares each case exactly once
 # ---------------------------------------------------------------------------