[WIP] Free energy fitting

espaloma/mm/implicit.py

@@ -0,0 +1,119 @@
+import torch
+
+torch.set_default_dtype(torch.float64)
+
+from espaloma.units import DISTANCE_UNIT, ENERGY_UNIT
+
+from simtk import unit
+
+distance_to_nm = (1.0 * DISTANCE_UNIT).value_in_unit(unit.nanometer)
+energy_from_kjmol = (1.0 * unit.kilojoule_per_mole).value_in_unit(ENERGY_UNIT)
+
+
+def step(x):
+    """return (x > 0)"""
+    return 1.0 * (x >= 0)
+
+
+def _gbsa_obc2_energy_omm(
+        distance_matrix,
+        radii, scales, charges,
+        alpha=0.8, beta=0.0, gamma=2.909125,
+        dielectric_offset=0.009,
+        surface_tension=28.3919551,
+        solute_dielectric=1.0,
+        solvent_dielectric=78.5,
+        probe_radius=0.14
+):
+    """
+    Assume everything is given in OpenMM units
+    ported from jax/numpy implementation here:
+    https://github.com/openforcefield/bayes-implicit-solvent/blob/067239fcbb8af28eb6310d702804887662692ec2/bayes_implicit_solvent/gb_models/jax_gb_models.py#L13-L60
+
+    with corrections and refinements by Yutong Zhao here
+    https://github.com/proteneer/timemachine/blob/417f4b0b1181b638935518532c78c380b03d7d19/timemachine/potentials/gbsa.py#L1-L111
+    """
+
+    N = len(charges)
+    eye = torch.eye(N, dtype=distance_matrix.dtype)
+
+    r = distance_matrix + eye
+    or1 = radii.reshape((N, 1)) - dielectric_offset
+    or2 = radii.reshape((1, N)) - dielectric_offset
+    sr2 = scales.reshape((1, N)) * or2
+
+    L = torch.max(or1, abs(r - sr2))
+    U = r + sr2
+
+    I = 1 / L - 1 / U + 0.25 * (r - sr2 ** 2 / r) * (
+            1 / (U ** 2) - 1 / (L ** 2)) + 0.5 * torch.log(
+        L / U) / r
+    # handle the interior case
+    I = torch.where(or1 < (sr2 - r), I + 2 * (1 / or1 - 1 / L), I)
+    I = step(r + sr2 - or1) * 0.5 * I  # note the extra 0.5 here
+    I -= torch.diag(torch.diag(I))
+    I = torch.sum(I, dim=1)
+
+    # okay, next compute born radii
+    offset_radius = radii - dielectric_offset
+
+    psi = I * offset_radius
+
+    psi_coefficient = alpha
+    psi2_coefficient = beta
+    psi3_coefficient = gamma
+
+    psi_term = (psi_coefficient * psi) - (psi2_coefficient * psi ** 2) + (
+            psi3_coefficient * psi ** 3)
+
+    B = 1 / (1 / offset_radius - torch.tanh(psi_term) / radii)
+
+    E = 0.0
+    # single particle
+    # ACE
+    E += torch.sum(
+        surface_tension * (radii + probe_radius) ** 2 * (radii / B) ** 6)
+
+    # on-diagonal
+    E += torch.sum(-0.5 * (
+            1 / solute_dielectric - 1 / solvent_dielectric) * charges ** 2 / B)
+
+    # particle pair
+    # note: np.outer --> torch.ger
+    f = torch.sqrt(r ** 2 + torch.ger(B, B) * torch.exp(
+        -r ** 2 / (4 * torch.ger(B, B))))
+    charge_products = torch.ger(charges, charges)
+
+    ixns = - (
+            1 / solute_dielectric - 1 / solvent_dielectric) * charge_products / f
+
+    E += torch.sum(torch.triu(ixns, diagonal=1))
+    return E  # E is in kJ/mol at this point
+
+
+def gbsa_obc2_energy(
+        distance_matrix_in_bohr,
+        radii_in_bohr, scales, charges,
+        alpha=0.8, beta=0.0, gamma=2.909125,
+        dielectric_offset=0.009,
+        surface_tension=28.3919551,
+        solute_dielectric=1.0,
+        solvent_dielectric=78.5,
+        probe_radius=0.14
+):
+    # convert distances and radii into units of nanometers before proceeding
+    distance_matrix = distance_matrix_in_bohr * distance_to_nm
+    radii = radii_in_bohr * distance_to_nm
+
+    E = _gbsa_obc2_energy_omm(
+        distance_matrix,
+        radii, scales, charges,
+        alpha, beta, gamma,
+        dielectric_offset=dielectric_offset,
+        surface_tension=surface_tension,
+        solute_dielectric=solute_dielectric,
+        solvent_dielectric=solvent_dielectric,
+        probe_radius=probe_radius,
+    )
+
+    return E * energy_from_kjmol  # return E in espaloma energy unit

espaloma/mm/tests/test_implicit.py

@@ -0,0 +1,69 @@
+
+
+import jax.numpy as np
+
+
+def step(x):
+    # return (x > 0)
+    return 1.0 * (x >= 0)
+
+
+def compute_OBC_energy_bayes_implicit(distance_matrix, radii, scales, charges,
+                                  offset=0.009, screening=138.935484,
+                                  surface_tension=28.3919551,
+                                  solvent_dielectric=78.5,
+                                  solute_dielectric=1.0,
+                                  ):
+    """From https://github.com/openforcefield/bayes-implicit-solvent/blob/46936da65ed93ed33f0f97362a1dea12f9170758/bayes_implicit_solvent/gb_models/jax_gb_models.py
+
+    in turn based on https://github.com/openmm/openmm/blob/master/platforms/reference/src/SimTKReference/ReferenceObc.cpp
+    """
+    N = len(radii)
+    # print(type(distance_matrix))
+    eye = np.eye(N, dtype=distance_matrix.dtype)
+    # print(type(eye))
+    r = distance_matrix + eye  # so I don't have divide-by-zero nonsense
+    or1 = radii.reshape((N, 1)) - offset
+    or2 = radii.reshape((1, N)) - offset
+    sr2 = scales.reshape((1, N)) * or2
+
+    L = np.maximum(or1, abs(r - sr2))
+    U = r + sr2
+    I = step(r + sr2 - or1) * 0.5 * (
+            1 / L - 1 / U + 0.25 * (r - sr2 ** 2 / r) * (
+                1 / (U ** 2) - 1 / (L ** 2)) + 0.5 * np.log(
+        L / U) / r)
+
+    I -= np.diag(np.diag(I))
+    I = np.sum(I, axis=1)
+
+    # okay, next compute born radii
+    offset_radius = radii - offset
+    psi = I * offset_radius
+    psi_coefficient = 0.8
+    psi2_coefficient = 0
+    psi3_coefficient = 2.909125
+
+    psi_term = (psi_coefficient * psi) + (psi2_coefficient * psi ** 2) + (
+                psi3_coefficient * psi ** 3)
+
+    B = 1 / (1 / offset_radius - np.tanh(psi_term) / radii)
+
+    # finally, compute the three energy terms
+    E = 0.0
+
+    # single particle
+    E += np.sum(surface_tension * (radii + 0.14) ** 2 * (radii / B) ** 6)
+    E += np.sum(-0.5 * screening * (
+                1 / solute_dielectric - 1 / solvent_dielectric) * charges ** 2 / B)
+
+    # particle pair
+    f = np.sqrt(
+        r ** 2 + np.outer(B, B) * np.exp(-r ** 2 / (4 * np.outer(B, B))))
+    charge_products = np.outer(charges, charges)
+
+    E += np.sum(np.triu(-screening * (
+                1 / solute_dielectric - 1 / solvent_dielectric) * charge_products / f,
+                        k=1))
+
+    return E

espaloma/redux/energy.py

@@ -6,30 +6,37 @@
 
 
 def compute_bonds(
-    xyz: torch.Tensor, params: ParameterizedSystem, indices: Indices
-) -> torch.Tensor:
+        xyz: torch.Tensor,
+        params: ParameterizedSystem,
+        indices: Indices) -> torch.Tensor:
     a, b = xyz[:, indices.bonds[:, 0]], xyz[:, indices.bonds[:, 1]]
     distance = esp.mm.geometry.distance(a, b)
     k, eq = params.bonds[:, 0], params.bonds[:, 1]
     return esp.mm.bond.harmonic_bond(distance, k, eq)
 
 
 def compute_angles(
-    xyz: torch.Tensor, params: ParameterizedSystem, indices: Indices
-) -> torch.Tensor:
-    a, b, c = (
-        xyz[:, indices.angles[:, 0]],
-        xyz[:, indices.angles[:, 1]],
-        xyz[:, indices.angles[:, 2]],
-    )
+        xyz: torch.Tensor,
+        params: ParameterizedSystem,
+        indices: Indices) -> torch.Tensor:
+    # TODO; be less verbose about this
+    a = xyz[:, indices.angles[:, 0]]
+    b = xyz[:, indices.angles[:, 1]]
+    c = xyz[:, indices.angles[:, 2]]
+
     angles = esp.mm.geometry.angle(a, b, c)
     k, eq = params.angles[:, 0], params.angles[:, 1]
     return esp.mm.angle.harmonic_angle(angles, k, eq)
 
 
 def compute_propers(
-    xyz: torch.Tensor, params: ParameterizedSystem, indices: Indices
-) -> torch.Tensor:
+        xyz: torch.Tensor,
+        params: ParameterizedSystem,
+        indices: Indices) -> torch.Tensor:
+    # it's possible there are no proper torsions in the system (e.g. h2o)
+    if len(indices.propers) == 0:
+        return torch.tensor([[0.0]])
+
     # TODO: reduce code duplication with compute_impropers
     a, b = xyz[:, indices.propers[:, 0]], xyz[:, indices.propers[:, 1]]
     c, d = xyz[:, indices.propers[:, 2]], xyz[:, indices.propers[:, 3]]
@@ -39,8 +46,13 @@ def compute_propers(
 
 
 def compute_impropers(
-    xyz: torch.Tensor, params: ParameterizedSystem, indices: Indices
-) -> torch.Tensor:
+        xyz: torch.Tensor,
+        params: ParameterizedSystem,
+        indices: Indices) -> torch.Tensor:
+    # it's possible there are no iproper torsions in the system (e.g. nh4)
+    if len(indices.impropers) == 0:
+        return torch.tensor([[0.0]])
+
     # TODO: reduce code duplication with compute_propers
     a, b = xyz[:, indices.impropers[:, 0]], xyz[:, indices.impropers[:, 1]]
     c, d = xyz[:, indices.impropers[:, 2]], xyz[:, indices.impropers[:, 3]]
@@ -50,13 +62,15 @@ def compute_impropers(
 
 
 def compute_valence_energy(
-    offmol: Molecule, xyz: torch.Tensor, params: ParameterizedSystem
-) -> torch.Tensor:
+        offmol: Molecule,
+        xyz: torch.Tensor,
+        params: ParameterizedSystem) -> torch.Tensor:
     indices = Indices(offmol)
-    harmonic_terms = compute_bonds(xyz, params, indices).sum(
-        1
-    ) + compute_angles(xyz, params, indices).sum(1)
-    torsion_terms = compute_propers(xyz, params, indices).sum(
-        1
-    ) + compute_impropers(xyz, params, indices).sum(1)
-    return harmonic_terms + torsion_terms
+
+    bonds = compute_bonds(xyz, params, indices).sum(1)
+    angles = compute_angles(xyz, params, indices).sum(1)
+    propers = compute_propers(xyz, params, indices).sum(1)
+    impropers = compute_impropers(xyz, params, indices).sum(1)
+    valence_energy = bonds + angles + propers + impropers
+
+    return valence_energy