torsion_parametrization

choderalab · Sep 12, 2021 · 36b88a8 · 36b88a8
1 parent adec67f
commit 36b88a8
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/espaloma/mm/functional.py b/espaloma/mm/functional.py
@@ -161,7 +161,13 @@ def periodic(
 
     k = k[:, None, :].repeat(1, x.shape[1], 1)
 
-    energy = (k * (1.0 + cos_n_theta_minus_phases)).sum(dim=-1)
+    # energy = (k * (1.0 + cos_n_theta_minus_phases)).sum(dim=-1)
+
+    energy = (
+        torch.nn.functional.relu(k) * (cos_n_theta_minus_phases + 1.0)
+       -torch.nn.functional.relu(0.0-k) * (cos_n_theta_minus_phases - 1.0)
+    ).sum(dim=-1)
+
 
     return energy