mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

hessian_ff/terms/dihedral.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+def _dihedral_angle_from_vectors(
+    v1: torch.Tensor,
+    v2: torch.Tensor,
+    v3: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """OpenMM-style signed dihedral from 3 bond-like vectors.
+
+    Returns
+    -------
+    phi, cross12, cross23
+    """
+    cross12 = torch.cross(v1, v2, dim=-1)
+    cross23 = torch.cross(v2, v3, dim=-1)
+
+    n12 = torch.linalg.norm(cross12, dim=-1).clamp_min(1.0e-12)
+    n23 = torch.linalg.norm(cross23, dim=-1).clamp_min(1.0e-12)
+    cos_phi = torch.sum(cross12 * cross23, dim=-1) / (n12 * n23)
+    cos_phi = torch.clamp(cos_phi, -1.0, 1.0)
+    phi = torch.acos(cos_phi)
+
+    sign = torch.where(
+        torch.sum(v1 * cross23, dim=-1) < 0.0,
+        -torch.ones_like(phi),
+        torch.ones_like(phi),
+    )
+    phi = phi * sign
+    return phi, cross12, cross23
+
+
+def _dihedral_angle(p0: torch.Tensor, p1: torch.Tensor, p2: torch.Tensor, p3: torch.Tensor) -> torch.Tensor:
+    """Return OpenMM-compatible signed dihedral angle (radians) for batched points."""
+    v1 = p0 - p1
+    v2 = p2 - p1
+    v3 = p2 - p3
+    phi, _, _ = _dihedral_angle_from_vectors(v1, v2, v3)
+    return phi
+
+
+def _accumulate_dihedral_forces(
+    force: torch.Tensor,
+    coords: torch.Tensor,
+    i: torch.Tensor,
+    j: torch.Tensor,
+    k: torch.Tensor,
+    l: torch.Tensor,
+    dE_dphi: torch.Tensor,
+) -> None:
+    """Accumulate torsion force for given dE/dphi following OpenMM reference formulas."""
+    if i.numel() == 0:
+        return
+
+    v1 = coords[i] - coords[j]  # A-B
+    v2 = coords[k] - coords[j]  # C-B
+    v3 = coords[k] - coords[l]  # C-D
+
+    cross12 = torch.cross(v1, v2, dim=-1)
+    cross23 = torch.cross(v2, v3, dim=-1)
+
+    norm_cross12_sq = torch.sum(cross12 * cross12, dim=-1).clamp_min(1.0e-24)
+    norm_cross23_sq = torch.sum(cross23 * cross23, dim=-1).clamp_min(1.0e-24)
+    norm_v2 = torch.linalg.norm(v2, dim=-1).clamp_min(1.0e-12)
+    norm_v2_sq = torch.sum(v2 * v2, dim=-1).clamp_min(1.0e-24)
+
+    f0 = (-dE_dphi * norm_v2) / norm_cross12_sq
+    f3 = (dE_dphi * norm_v2) / norm_cross23_sq
+    f1 = torch.sum(v1 * v2, dim=-1) / norm_v2_sq
+    f2 = torch.sum(v3 * v2, dim=-1) / norm_v2_sq
+
+    ff0 = f0.unsqueeze(-1) * cross12
+    ff3 = f3.unsqueeze(-1) * cross23
+    s = f1.unsqueeze(-1) * ff0 - f2.unsqueeze(-1) * ff3
+    ff1 = ff0 - s
+    ff2 = ff3 + s
+
+    force.index_add_(0, i, ff0)
+    force.index_add_(0, j, -ff1)
+    force.index_add_(0, k, -ff2)
+    force.index_add_(0, l, ff3)
+
+
+class DihedralTerm(nn.Module):
+    """AMBER torsion term: E = sum k (1 + cos(n*phi - phase))."""
+
+    def __init__(
+        self,
+        i: torch.Tensor,
+        j: torch.Tensor,
+        k: torch.Tensor,
+        l: torch.Tensor,
+        force: torch.Tensor,
+        period: torch.Tensor,
+        phase: torch.Tensor,
+    ):
+        super().__init__()
+        self.register_buffer("i", i.long())
+        self.register_buffer("j", j.long())
+        self.register_buffer("k", k.long())
+        self.register_buffer("l", l.long())
+        self.register_buffer("force", force)
+        self.register_buffer("period", period)
+        self.register_buffer("phase", phase)
+
+    def forward(self, coords: torch.Tensor) -> torch.Tensor:
+        if self.i.numel() == 0:
+            return coords.new_zeros(())
+        p0 = coords[self.i]
+        p1 = coords[self.j]
+        p2 = coords[self.k]
+        p3 = coords[self.l]
+        phi = _dihedral_angle(p0, p1, p2, p3)
+
+        # Periodicity can be negative in some prmtops; AMBER uses its absolute value.
+        n = torch.abs(self.period)
+        return torch.sum(self.force * (1.0 + torch.cos(n * phi - self.phase)))
+
+    def energy_force(self, coords: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Return torsion energy and analytical force."""
+        force = torch.zeros_like(coords)
+        if self.i.numel() == 0:
+            return coords.new_zeros(()), force
+
+        p0 = coords[self.i]
+        p1 = coords[self.j]
+        p2 = coords[self.k]
+        p3 = coords[self.l]
+        phi = _dihedral_angle(p0, p1, p2, p3)
+
+        n = torch.abs(self.period)
+        delta = n * phi - self.phase
+        energy = torch.sum(self.force * (1.0 + torch.cos(delta)))
+
+        # dE/dphi
+        dE_dphi = -self.force * n * torch.sin(delta)
+        _accumulate_dihedral_forces(force, coords, self.i, self.j, self.k, self.l, dE_dphi)
+        return energy, force

hessian_ff/terms/nonbonded.py

@@ -0,0 +1,209 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ..constants import COULOMB_K, validate_coords
+from ..native.nonbonded import (
+    nonbonded_energy_force_native,
+    nonbonded_energy_force_preparam_native,
+)
+
+
+@dataclass
+class NonbondedEnergies:
+    coulomb: torch.Tensor
+    lj: torch.Tensor
+    coulomb14: torch.Tensor
+    lj14: torch.Tensor
+
+
+class NonbondedTerm(nn.Module):
+    """No-PBC AMBER nonbonded term (native C++ backend only)."""
+
+    def __init__(
+        self,
+        natom: int,
+        charge: torch.Tensor,
+        atom_type: torch.Tensor,
+        lj_acoef: torch.Tensor,
+        lj_bcoef: torch.Tensor,
+        hb_acoef: torch.Tensor,
+        hb_bcoef: torch.Tensor,
+        nb_index: torch.Tensor,
+        pair_i: torch.Tensor,
+        pair_j: torch.Tensor,
+        pair14_i: torch.Tensor,
+        pair14_j: torch.Tensor,
+        pair14_inv_scee: torch.Tensor,
+        pair14_inv_scnb: torch.Tensor,
+        cpu_fast_kernel: bool = True,
+    ):
+        super().__init__()
+        self.natom = int(natom)
+        self.cpu_fast_kernel = bool(cpu_fast_kernel)
+
+        self.register_buffer("charge", charge)
+        self.register_buffer("atom_type", atom_type.long())
+        self.register_buffer("lj_acoef", lj_acoef)
+        self.register_buffer("lj_bcoef", lj_bcoef)
+        self.register_buffer("hb_acoef", hb_acoef)
+        self.register_buffer("hb_bcoef", hb_bcoef)
+        self.register_buffer("nb_index", nb_index.long())
+
+        self.register_buffer("pair_i", pair_i.long())
+        self.register_buffer("pair_j", pair_j.long())
+        self.register_buffer("pair14_i", pair14_i.long())
+        self.register_buffer("pair14_j", pair14_j.long())
+        self.register_buffer("pair14_inv_scee", pair14_inv_scee)
+        self.register_buffer("pair14_inv_scnb", pair14_inv_scnb)
+
+        # Precompute pair constants once so repeated coordinate updates on the
+        # same topology skip atom-type/table lookups.
+        p_coul, p_a12, p_b6, p_b10 = self._build_pair_params(
+            self.pair_i,
+            self.pair_j,
+            inv_scee=None,
+            inv_scnb=None,
+        )
+        p14_coul, p14_a12, p14_b6, p14_b10 = self._build_pair_params(
+            self.pair14_i,
+            self.pair14_j,
+            inv_scee=self.pair14_inv_scee,
+            inv_scnb=self.pair14_inv_scnb,
+        )
+        self.register_buffer("pair_coul_coeff", p_coul)
+        self.register_buffer("pair_a12_coeff", p_a12)
+        self.register_buffer("pair_b6_coeff", p_b6)
+        self.register_buffer("pair_b10_coeff", p_b10)
+        self.register_buffer("pair14_coul_coeff", p14_coul)
+        self.register_buffer("pair14_a12_coeff", p14_a12)
+        self.register_buffer("pair14_b6_coeff", p14_b6)
+        self.register_buffer("pair14_b10_coeff", p14_b10)
+
+    def _build_pair_params(
+        self,
+        pair_i: torch.Tensor,
+        pair_j: torch.Tensor,
+        inv_scee: torch.Tensor | None,
+        inv_scnb: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        dtype = self.charge.dtype
+        device = self.charge.device
+        n = int(pair_i.numel())
+        if n == 0:
+            z = torch.zeros((0,), dtype=dtype, device=device)
+            return z, z, z, z
+
+        ii = pair_i
+        jj = pair_j
+        qq = self.charge[ii] * self.charge[jj]
+        coul = qq * float(COULOMB_K)
+        if inv_scee is not None and inv_scee.numel() > 0:
+            coul = coul * inv_scee.to(dtype=dtype)
+
+        a12 = torch.zeros_like(coul)
+        b6 = torch.zeros_like(coul)
+        b10 = torch.zeros_like(coul)
+
+        ti = self.atom_type[ii]
+        tj = self.atom_type[jj]
+        raw = self.nb_index[ti, tj]
+
+        lj_mask = raw > 0
+        if bool(torch.any(lj_mask)):
+            lj_idx = raw[lj_mask] - 1
+            a = self.lj_acoef.index_select(0, lj_idx)
+            b = self.lj_bcoef.index_select(0, lj_idx)
+            if inv_scnb is not None and inv_scnb.numel() > 0:
+                s = inv_scnb[lj_mask].to(dtype=dtype)
+                a = a * s
+                b = b * s
+            a12.index_put_((lj_mask,), a)
+            b6.index_put_((lj_mask,), b)
+
+        hb_mask = raw < 0
+        if bool(torch.any(hb_mask)) and int(self.hb_acoef.numel()) > 0 and int(self.hb_bcoef.numel()) > 0:
+            hb_idx_full = (-raw[hb_mask]) - 1
+            valid = (
+                (hb_idx_full >= 0)
+                & (hb_idx_full < int(self.hb_acoef.numel()))
+                & (hb_idx_full < int(self.hb_bcoef.numel()))
+            )
+            if bool(torch.any(valid)):
+                hb_pos = torch.nonzero(hb_mask, as_tuple=False).reshape(-1)[valid]
+                hb_idx = hb_idx_full[valid]
+                a = self.hb_acoef.index_select(0, hb_idx)
+                b = self.hb_bcoef.index_select(0, hb_idx)
+                if inv_scnb is not None and inv_scnb.numel() > 0:
+                    s = inv_scnb[hb_pos].to(dtype=dtype)
+                    a = a * s
+                    b = b * s
+                a12.index_put_((hb_pos,), a)
+                b10.index_put_((hb_pos,), b)
+
+        return coul, a12, b6, b10
+
+    def _native_energy_force(
+        self,
+        coords: torch.Tensor,
+    ) -> tuple[NonbondedEnergies, torch.Tensor]:
+        validate_coords(
+            coords, self.natom, self.charge.dtype, self.charge.device,
+            label="coords (nonbonded)",
+        )
+        if not self.cpu_fast_kernel:
+            out = nonbonded_energy_force_preparam_native(
+                coords=coords,
+                pair_i=self.pair_i,
+                pair_j=self.pair_j,
+                pair_coul_coeff=self.pair_coul_coeff,
+                pair_a12_coeff=self.pair_a12_coeff,
+                pair_b6_coeff=self.pair_b6_coeff,
+                pair_b10_coeff=self.pair_b10_coeff,
+                pair14_i=self.pair14_i,
+                pair14_j=self.pair14_j,
+                pair14_coul_coeff=self.pair14_coul_coeff,
+                pair14_a12_coeff=self.pair14_a12_coeff,
+                pair14_b6_coeff=self.pair14_b6_coeff,
+                pair14_b10_coeff=self.pair14_b10_coeff,
+                cpu_fast=False,
+            )
+        else:
+            out = nonbonded_energy_force_native(
+                coords=coords,
+                charge=self.charge,
+                atom_type=self.atom_type,
+                lj_acoef=self.lj_acoef,
+                lj_bcoef=self.lj_bcoef,
+                hb_acoef=self.hb_acoef,
+                hb_bcoef=self.hb_bcoef,
+                nb_index=self.nb_index,
+                pair_i=self.pair_i,
+                pair_j=self.pair_j,
+                pair14_i=self.pair14_i,
+                pair14_j=self.pair14_j,
+                pair14_inv_scee=self.pair14_inv_scee,
+                pair14_inv_scnb=self.pair14_inv_scnb,
+                cpu_fast=self.cpu_fast_kernel,
+            )
+
+        e_coul, e_lj, e_coul14, e_lj14, force = out
+        return (
+            NonbondedEnergies(
+                coulomb=e_coul,
+                lj=e_lj,
+                coulomb14=e_coul14,
+                lj14=e_lj14,
+            ),
+            force,
+        )
+
+    def forward(self, coords: torch.Tensor) -> NonbondedEnergies:
+        return self._native_energy_force(coords)[0]
+
+    def energy_force(self, coords: torch.Tensor) -> tuple[NonbondedEnergies, torch.Tensor]:
+        """Return nonbonded energies and analytical force (native backend)."""
+        return self._native_energy_force(coords)

hessian_ff/tests/conftest.py

@@ -0,0 +1,75 @@
+"""Shared fixtures for hessian_ff tests."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+import torch
+
+from hessian_ff.forcefield import ForceFieldTorch
+from hessian_ff.loaders import load_coords, load_system
+
+DATA_DIR = Path(__file__).resolve().parent / "data" / "small"
+PRMTOP = DATA_DIR / "complex.parm7"
+COORDS = DATA_DIR / "complex.rst7"
+PDB = DATA_DIR / "complex.pdb"
+
+
+def _skip_no_data():
+    if not (PRMTOP.exists() and COORDS.exists()):
+        pytest.skip("hessian_ff test data (data/small/) not available")
+
+
+def _skip_no_openmm():
+    try:
+        from openmm import openmm as omm  # noqa: F401
+    except ImportError:
+        pytest.skip("OpenMM is not installed")
+
+
+@pytest.fixture(scope="session")
+def small_system():
+    """Load AmberSystem from small test data (float64, CPU)."""
+    _skip_no_data()
+    return load_system(PRMTOP, device="cpu").to(dtype=torch.float64)
+
+
+@pytest.fixture(scope="session")
+def small_ff(small_system):
+    """ForceFieldTorch for the small test system."""
+    return ForceFieldTorch(small_system, nonbonded_cpu_fast=True)
+
+
+@pytest.fixture(scope="session")
+def small_ff_autograd(small_system):
+    """ForceFieldTorch without nonbonded_cpu_fast (needed for autograd Hessian)."""
+    return ForceFieldTorch(small_system, nonbonded_cpu_fast=False)
+
+
+@pytest.fixture(scope="session")
+def small_coords(small_system):
+    """Coordinates tensor for the small test system (float64, CPU)."""
+    return load_coords(COORDS, natom=small_system.natom, device="cpu", dtype=torch.float64)
+
+
+@pytest.fixture(scope="session")
+def openmm_context_and_pos(small_system):
+    """Create an OpenMM context and base positions for the small system."""
+    _skip_no_openmm()
+    from openmm import app, openmm as omm, unit
+
+    prmtop = app.AmberPrmtopFile(str(PRMTOP))
+    system = prmtop.createSystem(
+        nonbondedMethod=app.NoCutoff, constraints=None, rigidWater=False,
+    )
+    integ = omm.VerletIntegrator(1.0 * unit.femtoseconds)
+    platform = omm.Platform.getPlatformByName("CPU")
+    context = omm.Context(system, integ, platform, {"Threads": "4"})
+
+    coords = load_coords(
+        COORDS, natom=small_system.natom, device="cpu", dtype=torch.float64,
+    )
+    pos_a = coords.detach().cpu().numpy()
+    context.setPositions(pos_a * unit.angstrom)
+    return context, pos_a