Lennard-Jones potential

docs/index.md

@@ -18,6 +18,7 @@
 ::: beignet.invert_quaternion
 ::: beignet.invert_rotation_matrix
 ::: beignet.invert_rotation_vector
+::: beignet.lennard_jones_potential
 ::: beignet.quaternion_identity
 ::: beignet.quaternion_magnitude
 ::: beignet.quaternion_mean

src/beignet/__init__.py

@@ -30,6 +30,7 @@
 from ._invert_rotation_matrix import invert_rotation_matrix
 from ._invert_rotation_vector import invert_rotation_vector
 from ._invert_transform import invert_transform
+from ._lennard_jones_potential import lennard_jones_potential
 from ._quaternion_identity import quaternion_identity
 from ._quaternion_magnitude import quaternion_magnitude
 from ._quaternion_mean import quaternion_mean
@@ -91,6 +92,7 @@
     "invert_rotation_matrix",
     "invert_rotation_vector",
     "invert_transform",
+    "lennard_jones_potential",
     "quaternion_identity",
     "quaternion_magnitude",
     "quaternion_mean",

src/beignet/_lennard_jones_potential.py

@@ -0,0 +1,33 @@
+from torch import Tensor
+
+
+def lennard_jones_potential(
+    input: Tensor,
+    sigma: float | Tensor,
+    epsilon: float | Tensor,
+) -> Tensor:
+    r"""
+    Lennard-Jones potential.
+
+    Parameters
+    ----------
+    input : Tensor, shape=(n, m)
+        Pairwise distances between particles.
+
+    sigma : float | Tensor, shape=(n, m)
+        Distance where the potential energy, :math:`V`, is zero.
+
+    epsilon : float | Tensor, shape=(n, m)
+        Depth of the potential well.
+
+    Returns
+    -------
+    output : Tensor, shape=(n, m)
+        Energies.
+    """
+    a = sigma / input
+
+    b = a**6.0
+    c = b**2.0
+
+    return 4.0 * epsilon * (c - b)

tests/beignet/test__lennard_jones_potential.py

@@ -0,0 +1,37 @@
+import beignet
+import torch
+import torch.testing
+
+
+def test__lennard_jones_potential():
+    x = beignet.lennard_jones_potential(
+        torch.tensor(
+            [
+                [1.0, 2.0],
+                [3.0, 4.0],
+            ],
+        ),
+        1.0,
+        1.0,
+    )
+
+    print(x)
+
+    torch.testing.assert_close(
+        beignet.lennard_jones_potential(
+            torch.tensor(
+                [
+                    [1.0, 2.0],
+                    [3.0, 4.0],
+                ],
+            ),
+            1.0,
+            1.0,
+        ),
+        torch.tensor(
+            [
+                [+0.0000000000000000, -0.0615234375000000],
+                [-0.0054794428870082, -0.0009763240814209],
+            ],
+        ),
+    )