Class ii

espaloma/mm/angle.py

@@ -30,8 +30,8 @@ def harmonic_angle(x, k, eq):
 
     # NOTE:
     # 0.25 because all angles are calculated twice
-    return 0.5 * esp.mm.functional.harmonic(x=x, k=k, eq=eq)
 
+    return 0.5 * esp.mm.functional.harmonic(x=x, k=k, eq=eq)
 
 def linear_mixture_angle(x, coefficients, phases):
     """ Angle energy with Linear basis function.

espaloma/mm/energy.py

@@ -40,10 +40,34 @@ def apply_angle(nodes, suffix=""):
             x=nodes.data["x"],
             k=nodes.data["k%s" % suffix],
             eq=nodes.data["eq%s" % suffix],
+        ),
+        "u_urey_bradley%s"
+        % suffix: esp.mm.angle.urey_bradley(
+            x_between=nodes.data["x_between"],
+            k_urey_bradley=nodes.data["k_urey_bradley"],
+            eq_urey_bradley=nodes.data["eq_urey_bradley"],
+        ),
+        "u_bond_bond%s"
+        % suffix: esp.mm.angle.bond_bond(
+            x_left=nodes.data["x_left"],
+            x_right=nodes.data["x_right"],
+            eq_left=nodes.data["eq_left"],
+            eq_right=nodes.data["eq_right"],
+            k_bond_bond=nodes.data["k_bond_bond"],
+        ),
+        "u_bond_angle%s"
+        % suffix: esp.mm.angle.bond_angle(
+            x_left=nodes.data["x_left"],
+            x_right=nodes.data["x_right"],
+            x_angle=nodes.data["x"],
+            eq_left=nodes.data["eq_left"],
+            eq_right=nodes.data["eq_right"],
+            eq_angle=nodes.data["eq"],
+            k_bond_angle_left=nodes.data["k_bond_angle_left"],
+            k_bond_angle_right=nodes.data["k_bond_angle_right"],
         )
     }
 
-
 def apply_angle_ii(nodes, suffix=""):
     return {
         "u_urey_bradley%s"
@@ -76,6 +100,23 @@ def apply_angle_ii(nodes, suffix=""):
 def apply_torsion_ii(nodes, suffix=""):
     """ Torsion energy in nodes. """
     return {
+        "u_angle_angle%s"
+        % suffix: esp.mm.torsion.angle_angle(
+            x_angle_left=nodes.data["x_angle_left"],
+            x_angle_right=nodes.data["x_angle_right"],
+            eq_angle_left=nodes.data["eq_angle_left"],
+            eq_angle_right=nodes.data["eq_angle_right"],
+            k_angle_angle=nodes.data["k_angle_angle"],
+        ),
+        "u_angle_angle_torsion%s"
+        % suffix: esp.mm.torsion.angle_angle_torsion(
+            x_angle_left=nodes.data["x_angle_left"],
+            x_angle_right=nodes.data["x_angle_right"],
+            eq_angle_left=nodes.data["eq_angle_left"],
+            eq_angle_right=nodes.data["eq_angle_right"],
+            x=nodes.data["x"],
+            k=nodes.data["k%s" % suffix],
+        ),
         "u_angle_angle%s"
         % suffix: esp.mm.torsion.angle_angle(
             x_angle_left=nodes.data["x_angle_left"],
@@ -316,6 +357,27 @@ def energy_in_graph(
             ntype="onefour",
         )
 
+    if class_ii is True:
+        g.apply_nodes(
+            lambda node: {
+                'u%s' % suffix:
+                node.data['u_urey_bradley%s' % suffix]\
+                + node.data['u_bond_bond%s' % suffix]\
+                + node.data['u_bond_angle%s' % suffix]
+            },
+            ntype='n3'
+        )
+
+        g.apply_nodes(
+            lambda node: {
+                'u%s' % suffix:
+                node.data['u_angle_angle%s' % suffix]\
+                + node.data['u_angle_angle_torsion%s' % suffix]\
+                + node.data['u_bond_torsion%s' % suffix]
+            },
+            ntype='n4'
+        )
+
     # sum up energy
     # bonded
     g.multi_update_all(

espaloma/mm/functional.py

@@ -285,7 +285,6 @@ def linear_mixture(x, coefficients, phases=[0.0, 1.0]):
 
     return u
 
-
 def harmonic_periodic_coupled(
         x_harmonic,
         x_periodic,

espaloma/mm/geometry.py

@@ -30,6 +30,9 @@ def _copy_src(edges, src=src, out=out):
 # =============================================================================
 def distance(x0, x1):
     """ Distance. """
+    # TODO:
+    # assertions / shape docstrings
+    # p=2 indicates Euclidian distance
     return torch.norm(x0 - x1, p=2, dim=-1)
 
 
@@ -55,7 +58,6 @@ def _dihedral(r0, r1):
     """ Dihedral between normal vectors. """
     return _angle(r0, r1)
 
-
 def dihedral(
     x0: torch.Tensor, x1: torch.Tensor, x2: torch.Tensor, x3: torch.Tensor
 ) -> torch.Tensor:
@@ -90,14 +92,19 @@ def dihedral(
     return theta
 
 
+    return torch.atan2(top, bottom)
+
 # =============================================================================
 # GEOMETRY IN HYPERNODES
 # =============================================================================
 def apply_bond(nodes):
     """ Bond length in nodes. """
-
-    return {"x": distance(x0=nodes.data["xyz0"], x1=nodes.data["xyz1"])}
-
+    return {
+        "x": distance(
+            x0=nodes.data["xyz0"],
+            x1=nodes.data["xyz1"]
+        ),
+    }
 
 def apply_angle(nodes):
     """ Angle values in nodes. """
@@ -121,7 +128,6 @@ def apply_angle(nodes):
         )
     }
 
-
 def apply_torsion(nodes):
     """ Torsion dihedrals in nodes. """
     return {
@@ -155,15 +161,12 @@ def apply_torsion(nodes):
         )
     }
 
-
-
 # =============================================================================
 # GEOMETRY IN GRAPH
 # =============================================================================
 # NOTE:
 # The following functions modify graphs in-place.
 
-
 def geometry_in_graph(g):
     """ Assign values to geometric entities in graphs.
 

espaloma/mm/tests/test_energy.py

@@ -50,7 +50,6 @@ def test_energy():
 
     esp.mm.energy.energy_in_graph(g, terms=["n2", "n3", "n4", "n4_improper"])
 
-
 # def test_energy_consistent():
 #     g = esp.Graph("c1ccccc1")
 #

espaloma/mm/torsion.py

@@ -7,11 +7,13 @@
 # =============================================================================
 # MODULE FUNCTIONS
 # =============================================================================
+
 def periodic_torsion(
     x, k, periodicity=list(range(1, 7)), phases=[0.0 for _ in range(6)]
 ):
     """ Periodic torsion potential
 
+
     Parameters
     ----------
     x : `torch.Tensor`, `shape = (batch_size, 1)`
@@ -123,3 +125,88 @@ def bond_torsion(
         k=k_right_torsion,
         eq=eq_right_torsion,
     )
+
+def angle_angle(
+        x_angle_left,
+        x_angle_right,
+        eq_angle_left,
+        eq_angle_right,
+        k_angle_angle,
+    ):
+    return esp.mm.functional.harmonic_harmonic_coupled(
+        x0=x_angle_left,
+        eq0=eq_angle_left,
+        x1=x_angle_right,
+        eq1=eq_angle_right,
+        k=k_angle_angle,
+    )
+
+def angle_torsion(
+        x_angle_left,
+        x_angle_right,
+        eq_angle_left,
+        eq_angle_right,
+        x,
+        k_angle_torsion_left,
+        k_angle_torsion_right,
+        periodicity=list(range(1,3)),
+    ):
+    return esp.mm.functional.harmonic_periodic_coupled(
+        x_harmonic=x_angle_left,
+        x_periodic=x,
+        k=k_angle_torsion_left,
+        eq=eq_angle_left,
+        periodicity=periodicity,
+    ) + esp.mm.functional.harmonic_periodic_coupled(
+        x_harmonic=x_angle_right,
+        x_periodic=x,
+        k=k_angle_torsion_right,
+        eq=eq_angle_right,
+        periodicity=periodicity,
+    )
+
+def angle_angle_torsion(
+        x_angle_left,
+        x_angle_right,
+        eq_angle_left,
+        eq_angle_right,
+        x,
+        k_angle_angle_torsion,
+    ):
+    return esp.mm.functional.harmonic_harmonic_periodic_coupled(
+        theta0=x_angle_left,
+        theta1=x_angle_right,
+        eq0=eq_angle_left,
+        eq1=eq_angle_right,
+        phi=x,
+        k=k_angle_angle_torsion,
+    )
+
+def bond_torsion(
+        x_bond_left,
+        x_bond_center,
+        x_bond_right,
+        x,
+        k_left_torsion,
+        k_center_torsion,
+        k_right_torsion,
+        eq_left_torsion,
+        eq_center_torsion,
+        eq_right_torsion,
+    ):
+    return esp.mm.functional.harmonic_periodic_coupled(
+        x_harmonic=x_bond_left,
+        x_periodic=x,
+        k=k_left_torsion,
+        eq=eq_left_torsion,
+    ) + esp.mm.functional.harmonic_periodic_coupled(
+        x_harmonic=x_bond_center,
+        x_periodic=x,
+        k=k_center_torsion,
+        eq=eq_center_torsion,
+    ) + esp.mm.functional.harmonic_periodic_coupled(
+        x_harmonic=x_bond_right,
+        x_periodic=x,
+        k=k_right_torsion,
+        eq=eq_right_torsion,
+    )