From 76e3aac4a3a119f2c61afd69f3da7c99d4fd19e3 Mon Sep 17 00:00:00 2001
From: Allen Goodman <allen.goodman@icloud.com>
Date: Wed, 12 Jun 2024 11:30:34 -0400
Subject: [PATCH] Lennard-Jones potential

---
 docs/index.md                                 |  1 +
 src/beignet/__init__.py                       |  2 +
 src/beignet/_lennard_jones_potential.py       | 33 +++++++++++++++++
 .../beignet/test__lennard_jones_potential.py  | 37 +++++++++++++++++++
 4 files changed, 73 insertions(+)
 create mode 100644 src/beignet/_lennard_jones_potential.py
 create mode 100644 tests/beignet/test__lennard_jones_potential.py

diff --git a/docs/index.md b/docs/index.md
index 53f1199f37..37b23d0d1a 100644
--- a/docs/index.md
+++ b/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
diff --git a/src/beignet/__init__.py b/src/beignet/__init__.py
index b8b82c85d7..89b8ffaa1a 100644
--- a/src/beignet/__init__.py
+++ b/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",
diff --git a/src/beignet/_lennard_jones_potential.py b/src/beignet/_lennard_jones_potential.py
new file mode 100644
index 0000000000..3057f0e06b
--- /dev/null
+++ b/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)
diff --git a/tests/beignet/test__lennard_jones_potential.py b/tests/beignet/test__lennard_jones_potential.py
new file mode 100644
index 0000000000..cec721f5f2
--- /dev/null
+++ 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],
+            ],
+        ),
+    )