From 877d58dfaa7e30868197157a5fb07b090c8740b5 Mon Sep 17 00:00:00 2001
From: maxdolan <maxdolan@eduroam-145-236.nomadic.bris.ac.uk>
Date: Mon, 17 Jun 2024 12:59:24 +0100
Subject: [PATCH 1/3] changed forcefields to be class-based

---
 pylj/forcefields.py            | 147 ++++++++++++++++++++-------------
 pylj/pairwise.py               |   8 +-
 pylj/tests/test_forcefields.py |  96 ++++++++++-----------
 3 files changed, 136 insertions(+), 115 deletions(-)

diff --git a/pylj/forcefields.py b/pylj/forcefields.py
index 1221906..a644ea0 100644
--- a/pylj/forcefields.py
+++ b/pylj/forcefields.py
@@ -2,8 +2,7 @@
 from numba import njit
 
 
-#Jit tag here had to be removed
-def lennard_jones(dr, constants, force=False):
+class lennard_jones(object):
     r"""Calculate the energy or force for a pair of particles using the
     Lennard-Jones (A/B variant) forcefield.
 
@@ -28,18 +27,22 @@ def lennard_jones(dr, constants, force=False):
     float: array_like
         The potential energy or force between the particles.
     """
-
-    if force:
-        return 12 * constants[0] * np.power(dr, -13) - (
-            6 * constants[1] * np.power(dr, -7))
-
-    else:
-        return constants[0] * np.power(dr, -12) - (
-            constants[1] * np.power(dr, -6))
-
-
-#Jit tag here had to be removed
-def lennard_jones_sigma_epsilon(dr, constants, force=False):
+    def __init__(self, dr, constants):
+        self.dr = dr
+        self.a = constants[0]
+        self.b = constants[1]
+        self.energy = self.update_energy()
+        self.force = self.update_force()
+
+    def update_energy(self):
+        self.energy = self.a * np.power(self.dr, -12) - (self.b * np.power(self.dr, -6))
+        return self.energy
+    
+    def update_force(self):
+        self.force = 12 * self.a * np.power(self.dr, -13) - (6 * self.b * np.power(self.dr, -7))
+        return self.force
+
+class lennard_jones_sigma_epsilon(object):
     r"""Calculate the energy or force for a pair of particles using the
     Lennard-Jones (sigma/epsilon variant) forcefield.
 
@@ -64,18 +67,24 @@ def lennard_jones_sigma_epsilon(dr, constants, force=False):
     float: array_like
         The potential energy or force between the particles.
     """
-
-    if force:
-        return 48 * constants[1] * np.power(constants[0], 12) * np.power(
-            dr, -13) - (24 * constants[1] * np.power(
-                constants[0], 6) * np.power(dr, -7))
-    else:
-        return 4 * constants[1] * np.power(constants[0], 12) * np.power(dr, -12) - (
-            4 * constants[1] * np.power(constants[0], 6) * np.power(dr, -6))
-
-
-#Jit tag here had to be removed
-def buckingham(dr, constants, force=False):
+    def __init__(self, dr, constants):
+        self.dr = dr
+        self.sigma = constants[0]
+        self.epsilon = constants[1]
+        self.energy = self.update_energy()
+        self.force = self.update_force()
+    
+    def update_energy(self):
+        self.energy = 4 * self.epsilon * np.power(self.sigma, 12) * np.power(self.dr, -12) - (
+                        4 * self.epsilon * np.power(self.sigma, 6) * np.power(self.dr, -6)) 
+        return self.energy 
+    
+    def update_force(self):
+        self.force = 48 * self.epsilon * np.power(self.sigma, 12) * np.power(
+            self.dr, -13) - (24 * self.epsilon * np.power(self.sigma, 6) * np.power(self.dr, -7))
+        return self.force
+
+class buckingham(object):
     r""" Calculate the energy or force for a pair of particles using the
     Buckingham forcefield.
 
@@ -100,16 +109,23 @@ def buckingham(dr, constants, force=False):
     float: array_like
         The potential energy or force between the particles.
     """
-    if force:
-        return constants[0] * constants[1] * np.exp(
-            - np.multiply(constants[1], dr)) - 6 * constants[2] / np.power(
-            dr, 7)
-    else:
-        return constants[0] * np.exp(
-            - np.multiply(constants[1], dr)) - constants[2] / np.power(dr, 6)
-
-
-def square_well(dr, constants, max_val=np.inf, force=False):
+    def __init__(self, dr, constants):
+        self.dr = dr
+        self.a = constants[0]
+        self.b = constants[1]
+        self.c = constants[2]
+        self.energy = self.update_energy()
+        self.force = self.update_force()
+    
+    def update_energy(self):
+        self.energy = self.a * np.exp(- np.multiply(self.b, self.dr)) - self.c / np.power(self.dr, 6)
+        return self.energy
+    
+    def update_force(self):
+        self.force = self.a * self.b * np.exp(- np.multiply(self.b, self.dr)) - 6 * self.c / np.power(self.dr, 7)
+        return self.force
+
+class square_well(object):
     r'''Calculate the energy or force for a pair of particles using a
     square well model.
 
@@ -146,26 +162,39 @@ def square_well(dr, constants, max_val=np.inf, force=False):
     float: array_like
         The potential energy between the particles.
     '''
-    if not isinstance(dr, np.ndarray):
-        if isinstance(dr, list):
-            dr = np.array(dr, dtype='float')
-        elif isinstance(dr, float):
-            dr = np.array([dr], dtype='float')
-
-    if force:
-        raise ValueError("Force is infinite at sigma <= dr < lambda * sigma")
-
-    else:
-        E = np.zeros_like(dr)
-        E[np.where(dr < constants[0])] = max_val
-        E[np.where(dr >= constants[2] * constants[1])] = 0
-
-        # apply mask for sigma <= dr < lambda * sigma
-        a = constants[1] <= dr
-        b = dr < constants[2] * constants[1]
-        E[np.where(a & b)] = -constants[0]
-
-        if len(E) == 1:
-            return float(E[0])
-        else:
-            return np.array(E, dtype='float')
+    def __init__(self, dr, constants, max_val=np.inf):
+
+        if not isinstance(dr, np.ndarray):
+            if isinstance(dr, list):
+                dr = np.array(dr, dtype='float')
+            elif isinstance(dr, float):
+                dr = np.array([dr], dtype='float')
+        
+        self.dr = dr
+        self.epsilon = constants[0]
+        self.sigma = constants[1]
+        self.lamda = constants[2] #Spelling as lamda not lambda to avoid calling python lambda function
+        self.max_val = max_val
+        self.energy = self.update_energy()
+        self.force = 0
+
+    def update_energy(self):
+            E = np.zeros_like(self.dr)
+            E = np.zeros_like(self.dr)
+            E[np.where(self.dr < self.epsilon)] = self.max_val
+            E[np.where(self.dr >= self.lamda * self.sigma)] = 0
+
+            # apply mask for sigma <= self.dr < lambda * sigma
+            a = self.sigma <= self.dr
+            b = self.dr < self.lamda * self.sigma
+            E[np.where(a & b)] = -self.epsilon
+
+            if len(E) == 1:
+                self.energy = float(E[0])
+            else:
+                self.energy = np.array(E, dtype='float')
+            
+            return self.energy
+        
+    def update_force(self):
+            raise ValueError("Force is infinite at sigma <= dr < lambda * sigma")
\ No newline at end of file
diff --git a/pylj/pairwise.py b/pylj/pairwise.py
index 8280a80..ce549f8 100644
--- a/pylj/pairwise.py
+++ b/pylj/pairwise.py
@@ -48,8 +48,8 @@ def compute_force(particles, box_length, cut_off, constants, forcefield, mass):
     distances, dx, dy = heavy.dist(
         particles["xposition"], particles["yposition"], box_length
     )
-    forces = forcefield(distances, constants, force=True)
-    energies = forcefield(distances, constants, force=False)
+    forces = forcefield(distances, constants).force
+    energies = forcefield(distances, constants).energy
     forces[np.where(distances > cut_off)] = 0.0
     energies[np.where(distances > cut_off)] = 0.0
     particles = update_accelerations(particles, forces, mass_kg, dx, dy, distances)
@@ -211,7 +211,7 @@ def compute_energy(particles, box_length, cut_off, constants, forcefield):
     distances, dx, dy = heavy.dist(
         particles["xposition"], particles["yposition"], box_length
     )
-    energies = forcefield(distances, constants, force=False)
+    energies = forcefield(distances, constants).energy
     energies[np.where(distances > cut_off)] = 0.0
     return distances, energies
 
@@ -248,7 +248,7 @@ def calculate_pressure(
     distances, dx, dy = heavy.dist(
         particles["xposition"], particles["yposition"], box_length
     )
-    forces = forcefield(distances, constants, force=True)
+    forces = forcefield(distances, constants).force
     forces[np.where(distances > cut_off)] = 0.0
     pres = np.sum(forces * distances)
     boltzmann_constant = 1.3806e-23  # joules / kelvin
diff --git a/pylj/tests/test_forcefields.py b/pylj/tests/test_forcefields.py
index d343495..59bc629 100644
--- a/pylj/tests/test_forcefields.py
+++ b/pylj/tests/test_forcefields.py
@@ -6,93 +6,85 @@
 class TestForcefields(unittest.TestCase):
     def test_lennard_jones_energy(self):
         a = forcefields.lennard_jones(2.0, [1.0, 1.0])
-        assert_almost_equal(a, -0.015380859)
+        assert_almost_equal(a.energy, -0.015380859)
         b = forcefields.lennard_jones([2.0, 4.5], [1.0, 1.0])
-        assert_almost_equal(b, [-0.015380859, -0.00012041])
+        assert_almost_equal(b.energy, [-0.015380859, -0.00012041])
         c = forcefields.lennard_jones([2.0, 4.5], [0.5, 3.5])
-        assert_almost_equal(c, [-0.05456543, -0.00042149])
+        assert_almost_equal(c.energy, [-0.05456543, -0.00042149])
         d = forcefields.lennard_jones([1.0, 1.5, 20.0], [5.0, 3.5])
-        assert_almost_equal(
-            d, [1.50000000, -0.268733500, -5.46874988e-08])
+        assert_almost_equal(d.energy, [1.50000000, -0.268733500, -5.46874988e-08])
         e = forcefields.lennard_jones([100.0, 200.0, 500.0], [100.0, 300.0])
-        assert_almost_equal(e, [0, 0, 0])
+        assert_almost_equal(e.energy, [0, 0, 0])
 
     def test_lennard_jones_force(self):
-        a = forcefields.lennard_jones(2.0, [1.0, 1.0], force=True)
-        assert_almost_equal(a, -0.045410156)
-        b = forcefields.lennard_jones([2.0, 4.0, 6.0], [1.0, 1.0], force=True)
-        assert_almost_equal(b, [-0.045410156, -3.66032124e-04, -2.14325517e-05])
-        c = forcefields.lennard_jones([2.0, 4.0, 6.0], [1.5, 4.0], force=True)
-        assert_almost_equal(c, [-0.185302734, -1.46457553e-03, -8.57325038e-05])
-        d = forcefields.lennard_jones(
-            [150.0, 300.0, 500.0], [200.0, 500.0], force=True)
-        assert_almost_equal(d, [-1.7558299e-12, -1.3717421e-14, -3.8400000e-16])
+        a = forcefields.lennard_jones(2.0, [1.0, 1.0])
+        assert_almost_equal(a.force, -0.045410156)
+        b = forcefields.lennard_jones([2.0, 4.0, 6.0], [1.0, 1.0])
+        assert_almost_equal(b.force, [-0.045410156, -3.66032124e-04, -2.14325517e-05])
+        c = forcefields.lennard_jones([2.0, 4.0, 6.0], [1.5, 4.0])
+        assert_almost_equal(c.force, [-0.185302734, -1.46457553e-03, -8.57325038e-05])
+        d = forcefields.lennard_jones([150.0, 300.0, 500.0], [200.0, 500.0])
+        assert_almost_equal(d.force, [-1.7558299e-12, -1.3717421e-14, -3.8400000e-16])
 
     def test_lennard_jones_sigma_epsilon_energy(self):
         a = forcefields.lennard_jones_sigma_epsilon(2.0, [1.0, 0.25])
-        assert_almost_equal(a, -0.015380859)
+        assert_almost_equal(a.energy, -0.015380859)
         b = forcefields.lennard_jones_sigma_epsilon([2.0, 1.0], [1.0, 0.25])
-        assert_almost_equal(b, [-0.015380859, 0])
+        assert_almost_equal(b.energy, [-0.015380859, 0])
         c = forcefields.lennard_jones_sigma_epsilon(
             [2.0, 1.0, 1.5], [0.5, 0.75])
-        assert_almost_equal(c, [-0.0007322, -0.0461425, -0.0041095])
+        assert_almost_equal(c.energy, [-0.0007322, -0.0461425, -0.0041095])
         d = forcefields.lennard_jones_sigma_epsilon(
             [400.0, 500.0, 600.0], [5e-10, 9e-9])
-        assert_almost_equal(d, [0, 0, 0])
+        assert_almost_equal(d.energy, [0, 0, 0])
 
     def test_lennard_jones_sigma_epsilon_force(self):
-        a = forcefields.lennard_jones_sigma_epsilon(
-            2.0, [1.0, 0.25], force=True)
-        assert_almost_equal(a, -0.045410156)
-        b = forcefields.lennard_jones_sigma_epsilon(
-            [2.0, 4.0], [1.0, 0.25], force=True)
-        assert_almost_equal(
-            b, [-0.0454102, -0.000366])
-        c = forcefields.lennard_jones_sigma_epsilon(
-            [3.0, 4.0], [3.0, 1.0], force=True)
-        assert_almost_equal(c, [8.0, -0.6877549])
+        a = forcefields.lennard_jones_sigma_epsilon(2.0, [1.0, 0.25])
+        assert_almost_equal(a.force, -0.045410156)
+        b = forcefields.lennard_jones_sigma_epsilon([2.0, 4.0], [1.0, 0.25])
+        assert_almost_equal(b.force, [-0.0454102, -0.000366])
+        c = forcefields.lennard_jones_sigma_epsilon([3.0, 4.0], [3.0, 1.0])
+        assert_almost_equal(c.force, [8.0, -0.6877549])
 
     def test_buckingham_energy(self):
         a = forcefields.buckingham(2.0, [1.0, 1.0, 1.0])
-        assert_almost_equal(a, 0.1197103832)
+        assert_almost_equal(a.energy, 0.1197103832)
         b = forcefields.buckingham([2.0], [1.0, 1.0, 1.0])
-        assert_almost_equal(b, 0.1197103832)
+        assert_almost_equal(b.energy, 0.1197103832)
         c = forcefields.buckingham([2.0, 4.0], [1.0, 1.0, 1.0])
-        assert_almost_equal(c, [0.1197103832, 0.0180715])
+        assert_almost_equal(c.energy, [0.1197103832, 0.0180715])
         d = forcefields.buckingham([2.0, 4.0, 5.0], [0.01, 0.01, 0.01])
-        assert_almost_equal(d, [0.0096457, 0.0096055, 0.0095117])
+        assert_almost_equal(d.energy, [0.0096457, 0.0096055, 0.0095117])
 
     def test_buckingham_force(self):
-        a = forcefields.buckingham(2.0, [1.0, 1.0, 1.0], force=True)
-        assert_almost_equal(a, 0.08846028324)
-        b = forcefields.buckingham([2.0], [1.0, 1.0, 1.0], force=True)
-        assert_almost_equal(b, 0.08846028324)
-        c = forcefields.buckingham(
-            [2.0, 1.0, 4.0], [1.5, 0.1, 2.0], force=True)
-        assert_almost_equal(c, [0.0290596, -11.8642744, 0.0998156])
+        a = forcefields.buckingham(2.0, [1.0, 1.0, 1.0])
+        assert_almost_equal(a.force, 0.08846028324)
+        b = forcefields.buckingham([2.0], [1.0, 1.0, 1.0])
+        assert_almost_equal(b.force, 0.08846028324)
+        c = forcefields.buckingham([2.0, 1.0, 4.0], [1.5, 0.1, 2.0])
+        assert_almost_equal(c.force, [0.0290596, -11.8642744, 0.0998156])
 
     def test_square_well_energy(self):
         a = forcefields.square_well(2.0, [1.0, 1.5, 2.0])
-        assert_equal(a, -1.0)
+        assert_equal(a.energy, -1.0)
         b = forcefields.square_well(0.5, [1.0, 2.0, 1.25])
-        assert_equal(b, float('inf'))
+        assert_equal(b.energy, float('inf'))
         c = forcefields.square_well(3.0, [0.5, 1.5, 1.25])
-        assert_equal(c, 0)
+        assert_equal(c.energy, 0)
         d = forcefields.square_well([2.0, 0.5], [1.0, 1.5, 2.0])
-        assert_equal(d, [-1.0, float('inf')])
+        assert_equal(d.energy, [-1.0, float('inf')])
         e = forcefields.square_well([3.0, 3.0, 0.25], [1.0, 1.5, 1.25])
-        assert_equal(e, [0, 0, float('inf')])
-        f = forcefields.square_well(
-            [3.0, 3.0, 0.25], [1.0, 1.5, 1.25], max_val=5000)
-        assert_equal(f, [0, 0, 5000])
+        assert_equal(e.energy, [0, 0, float('inf')])
+        f = forcefields.square_well([3.0, 3.0, 0.25], [1.0, 1.5, 1.25], max_val=5000)
+        assert_equal(f.energy, [0, 0, 5000])
 
     def test_square_well_force(self):
+        a = forcefields.square_well(2.0, [1.0, 1.5, 2.0])
         with self.assertRaises(ValueError):
-            forcefields.square_well(
-                2.0, [1.0, 1.5, 2.0], force=True)
+            a.update_force()
+        b = forcefields.square_well([2.0], [1.0, 1.5, 2.0])
         with self.assertRaises(ValueError):
-            forcefields.square_well(
-                [2.0], [1.0, 1.5, 2.0], force=True)
+            b.update_force()
 
 
 if __name__ == '__main__':

From 449abc2e24d2953eae940fd542528e047eedb559 Mon Sep 17 00:00:00 2001
From: maxdolan <maxdolan@eduroam-145-236.nomadic.bris.ac.uk>
Date: Mon, 17 Jun 2024 13:19:36 +0100
Subject: [PATCH 2/3] updating forcefields docstrings

---
 pylj/forcefields.py | 196 ++++++++++++++++++++++++++------------------
 1 file changed, 118 insertions(+), 78 deletions(-)

diff --git a/pylj/forcefields.py b/pylj/forcefields.py
index a644ea0..b4ea541 100644
--- a/pylj/forcefields.py
+++ b/pylj/forcefields.py
@@ -6,12 +6,6 @@ class lennard_jones(object):
     r"""Calculate the energy or force for a pair of particles using the
     Lennard-Jones (A/B variant) forcefield.
 
-    .. math::
-        E = \frac{A}{dr^{12}} - \frac{B}{dr^6}
-
-    .. math::
-        f = \frac{12A}{dr^{13}} - \frac{6B}{dr^7}
-
     Parameters
     ----------
     dr: float, array_like
@@ -19,13 +13,6 @@ class lennard_jones(object):
     constants: float, array_like
         An array of length two consisting of the A and B parameters for the
         12-6 Lennard-Jones function
-    force: bool (optional)
-        If true, the negative first derivative will be found.
-
-    Returns
-    -------
-    float: array_like
-        The potential energy or force between the particles.
     """
     def __init__(self, dr, constants):
         self.dr = dr
@@ -35,23 +22,41 @@ def __init__(self, dr, constants):
         self.force = self.update_force()
 
     def update_energy(self):
+        r"""Calculate the energy for a pair of particles using the
+        Lennard-Jones (A/B variant) forcefield.
+
+        .. math::
+            E = \frac{A}{dr^{12}} - \frac{B}{dr^6}
+
+        Returns
+        -------
+        float: array_like
+        The potential energy between the particles.
+        """
         self.energy = self.a * np.power(self.dr, -12) - (self.b * np.power(self.dr, -6))
         return self.energy
     
     def update_force(self):
+        r"""Calculate the force for a pair of particles using the
+        Lennard-Jones (A/B variant) forcefield.
+
+        .. math::
+            f = \frac{12A}{dr^{13}} - \frac{6B}{dr^7}
+
+        Returns
+        -------
+        float: array_like
+        The force between the particles.
+        """
         self.force = 12 * self.a * np.power(self.dr, -13) - (6 * self.b * np.power(self.dr, -7))
         return self.force
 
+
+
 class lennard_jones_sigma_epsilon(object):
     r"""Calculate the energy or force for a pair of particles using the
     Lennard-Jones (sigma/epsilon variant) forcefield.
 
-    .. math::
-        E = \frac{4e*a^{12}}{dr^{12}} - \frac{4e*a^{6}}{dr^6}
-
-    .. math::
-        f = \frac{48e*a^{12}}{dr^{13}} - \frac{24e*a^{6}}{dr^7}
-
     Parameters
     ----------
     dr: float, array_like
@@ -59,13 +64,6 @@ class lennard_jones_sigma_epsilon(object):
     constants: float, array_like
         An array of length two consisting of the sigma (a) and epsilon (e)
         parameters for the 12-6 Lennard-Jones function
-    force: bool (optional)
-        If true, the negative first derivative will be found.
-
-    Returns
-    -------
-    float: array_like
-        The potential energy or force between the particles.
     """
     def __init__(self, dr, constants):
         self.dr = dr
@@ -75,25 +73,43 @@ def __init__(self, dr, constants):
         self.force = self.update_force()
     
     def update_energy(self):
+        r"""Calculate the energy for a pair of particles using the
+        Lennard-Jones (sigma/epsilon variant) forcefield.
+
+        .. math::
+            E = \frac{4e*a^{12}}{dr^{12}} - \frac{4e*a^{6}}{dr^6}
+
+        Returns
+        -------
+        float: array_like
+        The potential energy between the particles.
+        """
         self.energy = 4 * self.epsilon * np.power(self.sigma, 12) * np.power(self.dr, -12) - (
                         4 * self.epsilon * np.power(self.sigma, 6) * np.power(self.dr, -6)) 
         return self.energy 
     
     def update_force(self):
+        r"""Calculate the force for a pair of particles using the
+        Lennard-Jones (sigma/epsilon variant) forcefield.
+
+        .. math::
+            f = \frac{48e*a^{12}}{dr^{13}} - \frac{24e*a^{6}}{dr^7}
+
+        Returns
+        -------
+        float: array_like
+        The force between the particles.
+        """
         self.force = 48 * self.epsilon * np.power(self.sigma, 12) * np.power(
             self.dr, -13) - (24 * self.epsilon * np.power(self.sigma, 6) * np.power(self.dr, -7))
         return self.force
 
+
+
 class buckingham(object):
     r""" Calculate the energy or force for a pair of particles using the
     Buckingham forcefield.
 
-    .. math::
-        E = Ae^{(-Bdr)} - \frac{C}{dr^6}
-
-    .. math::
-        f = ABe^{(-Bdr)} - \frac{6C}{dr^7}
-
     Parameters
     ----------
     dr: float, array_like
@@ -101,13 +117,6 @@ class buckingham(object):
     constants: float, array_like
         An array of length three consisting of the A, B and C parameters for
         the Buckingham function.
-    force: bool (optional)
-        If true, the negative first derivative will be found.
-
-    Returns
-    -------
-    float: array_like
-        The potential energy or force between the particles.
     """
     def __init__(self, dr, constants):
         self.dr = dr
@@ -118,33 +127,41 @@ def __init__(self, dr, constants):
         self.force = self.update_force()
     
     def update_energy(self):
+        r"""Calculate the energy for a pair of particles using the
+        Buckingham forcefield.
+
+        .. math::
+            E = Ae^{(-Bdr)} - \frac{C}{dr^6}
+
+        Returns
+        -------
+        float: array_like
+        The potential energy between the particles.
+        """
         self.energy = self.a * np.exp(- np.multiply(self.b, self.dr)) - self.c / np.power(self.dr, 6)
         return self.energy
     
     def update_force(self):
+        r"""Calculate the force for a pair of particles using the
+        Buckingham forcefield.
+
+        .. math::
+            f = ABe^{(-Bdr)} - \frac{6C}{dr^7}
+
+        Returns
+        -------
+        float: array_like
+        The force between the particles.
+        """
         self.force = self.a * self.b * np.exp(- np.multiply(self.b, self.dr)) - 6 * self.c / np.power(self.dr, 7)
         return self.force
 
+
+
 class square_well(object):
     r'''Calculate the energy or force for a pair of particles using a
     square well model.
 
-    .. math::
-        E = {
-        if dr < sigma:
-            E = max_val
-        elif sigma <= dr < lambda * sigma:
-            E = -epsilon
-        elif r >= lambda * sigma:
-            E = 0
-        }
-    .. math::
-        f = {
-        if sigma <= dr < lambda * sigma:
-            f = inf
-        else:
-            f = 0
-        }
     Parameters
     ----------
     dr: float, array_like
@@ -154,13 +171,6 @@ class square_well(object):
         parameters for the square well model.
     max_val: int (optional)
         Upper bound for values in square well - replaces usual infinite values
-    force: bool (optional)
-        If true, the negative first derivative will be found.
-
-    Returns
-    -------
-    float: array_like
-        The potential energy between the particles.
     '''
     def __init__(self, dr, constants, max_val=np.inf):
 
@@ -179,22 +189,52 @@ def __init__(self, dr, constants, max_val=np.inf):
         self.force = 0
 
     def update_energy(self):
-            E = np.zeros_like(self.dr)
-            E = np.zeros_like(self.dr)
-            E[np.where(self.dr < self.epsilon)] = self.max_val
-            E[np.where(self.dr >= self.lamda * self.sigma)] = 0
-
-            # apply mask for sigma <= self.dr < lambda * sigma
-            a = self.sigma <= self.dr
-            b = self.dr < self.lamda * self.sigma
-            E[np.where(a & b)] = -self.epsilon
-
-            if len(E) == 1:
-                self.energy = float(E[0])
-            else:
-                self.energy = np.array(E, dtype='float')
+        r'''Calculate the energy for a pair of particles using a
+        square well model.
+
+        .. math::
+            E = {
+            if dr < sigma:
+                E = max_val
+            elif sigma <= dr < lambda * sigma:
+                E = -epsilon
+            elif r >= lambda * sigma:
+                E = 0
+            }
+
+        Returns
+        -------
+        float: array_like
+            The potential energy between the particles.
+        '''
+        E = np.zeros_like(self.dr)
+        E = np.zeros_like(self.dr)
+        E[np.where(self.dr < self.epsilon)] = self.max_val
+        E[np.where(self.dr >= self.lamda * self.sigma)] = 0
+
+        # apply mask for sigma <= self.dr < lambda * sigma
+        a = self.sigma <= self.dr
+        b = self.dr < self.lamda * self.sigma
+        E[np.where(a & b)] = -self.epsilon
+
+        if len(E) == 1:
+            self.energy = float(E[0])
+        else:
+            self.energy = np.array(E, dtype='float')
             
-            return self.energy
+        return self.energy
         
     def update_force(self):
-            raise ValueError("Force is infinite at sigma <= dr < lambda * sigma")
\ No newline at end of file
+        r'''The force of a pair of particles using a square well model is given by:
+
+        .. math::
+        f = {
+        if sigma <= dr < lambda * sigma:
+            f = inf
+        else:
+            f = 0
+        }
+
+        Therefore the force here will always be infinite, and therefore not possible to simulate
+        '''
+        raise ValueError("Force is infinite at sigma <= dr < lambda * sigma")
\ No newline at end of file

From f5061e460939e7c5017c5bfa1f7214af58474466 Mon Sep 17 00:00:00 2001
From: maxdolan <maxdolan@eduroam-145-236.nomadic.bris.ac.uk>
Date: Mon, 17 Jun 2024 15:20:07 +0100
Subject: [PATCH 3/3] adjusting energy/force methods in classes

---
 pylj/forcefields.py            | 119 +++++++++++++++++--------------
 pylj/pairwise.py               |  11 +--
 pylj/tests/test_forcefields.py | 126 ++++++++++++++++-----------------
 3 files changed, 135 insertions(+), 121 deletions(-)

diff --git a/pylj/forcefields.py b/pylj/forcefields.py
index b4ea541..bdb2e51 100644
--- a/pylj/forcefields.py
+++ b/pylj/forcefields.py
@@ -8,47 +8,51 @@ class lennard_jones(object):
 
     Parameters
     ----------
-    dr: float, array_like
-        The distances between the all pairs of particles.
     constants: float, array_like
         An array of length two consisting of the A and B parameters for the
         12-6 Lennard-Jones function
+
     """
-    def __init__(self, dr, constants):
-        self.dr = dr
+    def __init__(self, constants):
         self.a = constants[0]
         self.b = constants[1]
-        self.energy = self.update_energy()
-        self.force = self.update_force()
 
-    def update_energy(self):
+    def energy(self, dr ):
         r"""Calculate the energy for a pair of particles using the
         Lennard-Jones (A/B variant) forcefield.
 
         .. math::
             E = \frac{A}{dr^{12}} - \frac{B}{dr^6}
 
+        Attributes:
+        ----------
+        dr (float): The distance between particles.
+
         Returns
         -------
         float: array_like
         The potential energy between the particles.
         """
-        self.energy = self.a * np.power(self.dr, -12) - (self.b * np.power(self.dr, -6))
+        self.energy = self.a * np.power(dr, -12) - (self.b * np.power(dr, -6))
         return self.energy
     
-    def update_force(self):
+    def force(self, dr):
         r"""Calculate the force for a pair of particles using the
         Lennard-Jones (A/B variant) forcefield.
 
         .. math::
             f = \frac{12A}{dr^{13}} - \frac{6B}{dr^7}
 
+        Attributes:
+        ----------
+        dr (float): The distance between particles.
+        
         Returns
         -------
         float: array_like
         The force between the particles.
         """
-        self.force = 12 * self.a * np.power(self.dr, -13) - (6 * self.b * np.power(self.dr, -7))
+        self.force = 12 * self.a * np.power(dr, -13) - (6 * self.b * np.power(dr, -7))
         return self.force
 
 
@@ -59,49 +63,53 @@ class lennard_jones_sigma_epsilon(object):
 
     Parameters
     ----------
-    dr: float, array_like
-        The distances between the all pairs of particles.
     constants: float, array_like
         An array of length two consisting of the sigma (a) and epsilon (e)
         parameters for the 12-6 Lennard-Jones function
+
     """
-    def __init__(self, dr, constants):
-        self.dr = dr
+    def __init__(self, constants):
         self.sigma = constants[0]
         self.epsilon = constants[1]
-        self.energy = self.update_energy()
-        self.force = self.update_force()
     
-    def update_energy(self):
+    def energy(self, dr):
         r"""Calculate the energy for a pair of particles using the
         Lennard-Jones (sigma/epsilon variant) forcefield.
 
         .. math::
             E = \frac{4e*a^{12}}{dr^{12}} - \frac{4e*a^{6}}{dr^6}
 
+        Attributes:
+        ----------
+        dr (float): The distance between particles.
+ 
         Returns
         -------
         float: array_like
         The potential energy between the particles.
         """
-        self.energy = 4 * self.epsilon * np.power(self.sigma, 12) * np.power(self.dr, -12) - (
-                        4 * self.epsilon * np.power(self.sigma, 6) * np.power(self.dr, -6)) 
+        self.energy = 4 * self.epsilon * np.power(self.sigma, 12) * np.power(dr, -12) - (
+                        4 * self.epsilon * np.power(self.sigma, 6) * np.power(dr, -6)) 
         return self.energy 
     
-    def update_force(self):
+    def force(self, dr):
         r"""Calculate the force for a pair of particles using the
         Lennard-Jones (sigma/epsilon variant) forcefield.
 
         .. math::
             f = \frac{48e*a^{12}}{dr^{13}} - \frac{24e*a^{6}}{dr^7}
 
+        Attributes:
+        ----------
+        dr (float): The distance between particles.
+ 
         Returns
         -------
         float: array_like
         The force between the particles.
         """
         self.force = 48 * self.epsilon * np.power(self.sigma, 12) * np.power(
-            self.dr, -13) - (24 * self.epsilon * np.power(self.sigma, 6) * np.power(self.dr, -7))
+            dr, -13) - (24 * self.epsilon * np.power(self.sigma, 6) * np.power(dr, -7))
         return self.force
 
 
@@ -112,48 +120,52 @@ class buckingham(object):
 
     Parameters
     ----------
-    dr: float, array_like
-        The distances between all the pairs of particles.
     constants: float, array_like
         An array of length three consisting of the A, B and C parameters for
         the Buckingham function.
+
     """
-    def __init__(self, dr, constants):
-        self.dr = dr
+    def __init__(self, constants):
         self.a = constants[0]
         self.b = constants[1]
         self.c = constants[2]
-        self.energy = self.update_energy()
-        self.force = self.update_force()
     
-    def update_energy(self):
+    def energy(self, dr):
         r"""Calculate the energy for a pair of particles using the
         Buckingham forcefield.
 
         .. math::
             E = Ae^{(-Bdr)} - \frac{C}{dr^6}
 
+        Attributes:
+        ----------
+        dr (float): The distance between particles.
+ 
         Returns
         -------
         float: array_like
         The potential energy between the particles.
         """
-        self.energy = self.a * np.exp(- np.multiply(self.b, self.dr)) - self.c / np.power(self.dr, 6)
+        self.energy = self.a * np.exp(- np.multiply(self.b, dr)) - self.c / np.power(dr, 6)
         return self.energy
     
-    def update_force(self):
+    def force(self, dr):
         r"""Calculate the force for a pair of particles using the
         Buckingham forcefield.
 
         .. math::
             f = ABe^{(-Bdr)} - \frac{6C}{dr^7}
 
+        Attributes:
+        ----------
+        dr (float): The distance between particles.
+ 
         Returns
         -------
         float: array_like
         The force between the particles.
         """
-        self.force = self.a * self.b * np.exp(- np.multiply(self.b, self.dr)) - 6 * self.c / np.power(self.dr, 7)
+        self.force = self.a * self.b * np.exp(- np.multiply(self.b, dr)) - 6 * self.c / np.power(dr, 7)
         return self.force
 
 
@@ -164,31 +176,21 @@ class square_well(object):
 
     Parameters
     ----------
-    dr: float, array_like
-        The distances between all the pairs of particles.
     constants: float, array_like
         An array of length three consisting of the epsilon, sigma, and lambda
         parameters for the square well model.
     max_val: int (optional)
         Upper bound for values in square well - replaces usual infinite values
+        
     '''
-    def __init__(self, dr, constants, max_val=np.inf):
+    def __init__(self, constants, max_val=np.inf):
 
-        if not isinstance(dr, np.ndarray):
-            if isinstance(dr, list):
-                dr = np.array(dr, dtype='float')
-            elif isinstance(dr, float):
-                dr = np.array([dr], dtype='float')
-        
-        self.dr = dr
         self.epsilon = constants[0]
         self.sigma = constants[1]
         self.lamda = constants[2] #Spelling as lamda not lambda to avoid calling python lambda function
         self.max_val = max_val
-        self.energy = self.update_energy()
-        self.force = 0
 
-    def update_energy(self):
+    def energy(self, dr):
         r'''Calculate the energy for a pair of particles using a
         square well model.
 
@@ -202,19 +204,30 @@ def update_energy(self):
                 E = 0
             }
 
+        Attributes:
+        ----------
+        dr (float): The distance between particles.
+ 
         Returns
         -------
         float: array_like
             The potential energy between the particles.
         '''
-        E = np.zeros_like(self.dr)
-        E = np.zeros_like(self.dr)
-        E[np.where(self.dr < self.epsilon)] = self.max_val
-        E[np.where(self.dr >= self.lamda * self.sigma)] = 0
-
-        # apply mask for sigma <= self.dr < lambda * sigma
-        a = self.sigma <= self.dr
-        b = self.dr < self.lamda * self.sigma
+
+        if not isinstance(dr, np.ndarray):
+            if isinstance(dr, list):
+                dr = np.array(dr, dtype='float')
+            elif isinstance(dr, float):
+                dr = np.array([dr], dtype='float')
+
+        E = np.zeros_like(dr)
+        E = np.zeros_like(dr)
+        E[np.where(dr < self.epsilon)] = self.max_val
+        E[np.where(dr >= self.lamda * self.sigma)] = 0
+
+        # apply mask for sigma <= dr < lambda * sigma
+        a = self.sigma <= dr
+        b = dr < self.lamda * self.sigma
         E[np.where(a & b)] = -self.epsilon
 
         if len(E) == 1:
@@ -224,7 +237,7 @@ def update_energy(self):
             
         return self.energy
         
-    def update_force(self):
+    def force(self):
         r'''The force of a pair of particles using a square well model is given by:
 
         .. math::
diff --git a/pylj/pairwise.py b/pylj/pairwise.py
index ce549f8..fd0e26a 100644
--- a/pylj/pairwise.py
+++ b/pylj/pairwise.py
@@ -48,8 +48,9 @@ def compute_force(particles, box_length, cut_off, constants, forcefield, mass):
     distances, dx, dy = heavy.dist(
         particles["xposition"], particles["yposition"], box_length
     )
-    forces = forcefield(distances, constants).force
-    energies = forcefield(distances, constants).energy
+    ff = forcefield(constants)
+    forces = ff.force(distances)
+    energies = ff.energy(distances)
     forces[np.where(distances > cut_off)] = 0.0
     energies[np.where(distances > cut_off)] = 0.0
     particles = update_accelerations(particles, forces, mass_kg, dx, dy, distances)
@@ -211,7 +212,8 @@ def compute_energy(particles, box_length, cut_off, constants, forcefield):
     distances, dx, dy = heavy.dist(
         particles["xposition"], particles["yposition"], box_length
     )
-    energies = forcefield(distances, constants).energy
+    ff = forcefield(constants)
+    energies = ff.energy(distances)
     energies[np.where(distances > cut_off)] = 0.0
     return distances, energies
 
@@ -248,7 +250,8 @@ def calculate_pressure(
     distances, dx, dy = heavy.dist(
         particles["xposition"], particles["yposition"], box_length
     )
-    forces = forcefield(distances, constants).force
+    ff = forcefield(constants)
+    forces = ff.force(distances)
     forces[np.where(distances > cut_off)] = 0.0
     pres = np.sum(forces * distances)
     boltzmann_constant = 1.3806e-23  # joules / kelvin
diff --git a/pylj/tests/test_forcefields.py b/pylj/tests/test_forcefields.py
index 59bc629..a1376d8 100644
--- a/pylj/tests/test_forcefields.py
+++ b/pylj/tests/test_forcefields.py
@@ -5,86 +5,84 @@
 
 class TestForcefields(unittest.TestCase):
     def test_lennard_jones_energy(self):
-        a = forcefields.lennard_jones(2.0, [1.0, 1.0])
-        assert_almost_equal(a.energy, -0.015380859)
-        b = forcefields.lennard_jones([2.0, 4.5], [1.0, 1.0])
-        assert_almost_equal(b.energy, [-0.015380859, -0.00012041])
-        c = forcefields.lennard_jones([2.0, 4.5], [0.5, 3.5])
-        assert_almost_equal(c.energy, [-0.05456543, -0.00042149])
-        d = forcefields.lennard_jones([1.0, 1.5, 20.0], [5.0, 3.5])
-        assert_almost_equal(d.energy, [1.50000000, -0.268733500, -5.46874988e-08])
-        e = forcefields.lennard_jones([100.0, 200.0, 500.0], [100.0, 300.0])
-        assert_almost_equal(e.energy, [0, 0, 0])
+        a = forcefields.lennard_jones([1.0, 1.0])
+        assert_almost_equal(a.energy(2.0), -0.015380859)
+        b = forcefields.lennard_jones([1.0, 1.0])
+        assert_almost_equal(b.energy([2.0, 4.5],), [-0.015380859, -0.00012041])
+        c = forcefields.lennard_jones([0.5, 3.5])
+        assert_almost_equal(c.energy([2.0, 4.5],), [-0.05456543, -0.00042149])
+        d = forcefields.lennard_jones( [5.0, 3.5])
+        assert_almost_equal(d.energy([1.0, 1.5, 20.0]), [1.50000000, -0.268733500, -5.46874988e-08])
+        e = forcefields.lennard_jones([100.0, 300.0])
+        assert_almost_equal(e.energy([100.0, 200.0, 500.0]), [0, 0, 0])
 
     def test_lennard_jones_force(self):
-        a = forcefields.lennard_jones(2.0, [1.0, 1.0])
-        assert_almost_equal(a.force, -0.045410156)
-        b = forcefields.lennard_jones([2.0, 4.0, 6.0], [1.0, 1.0])
-        assert_almost_equal(b.force, [-0.045410156, -3.66032124e-04, -2.14325517e-05])
-        c = forcefields.lennard_jones([2.0, 4.0, 6.0], [1.5, 4.0])
-        assert_almost_equal(c.force, [-0.185302734, -1.46457553e-03, -8.57325038e-05])
-        d = forcefields.lennard_jones([150.0, 300.0, 500.0], [200.0, 500.0])
-        assert_almost_equal(d.force, [-1.7558299e-12, -1.3717421e-14, -3.8400000e-16])
+        a = forcefields.lennard_jones([1.0, 1.0])
+        assert_almost_equal(a.force(2.0), -0.045410156)
+        b = forcefields.lennard_jones([1.0, 1.0])
+        assert_almost_equal(b.force([2.0, 4.0, 6.0]), [-0.045410156, -3.66032124e-04, -2.14325517e-05])
+        c = forcefields.lennard_jones([1.5, 4.0])
+        assert_almost_equal(c.force([2.0, 4.0, 6.0]), [-0.185302734, -1.46457553e-03, -8.57325038e-05])
+        d = forcefields.lennard_jones([200.0, 500.0])
+        assert_almost_equal(d.force([150.0, 300.0, 500.0]), [-1.7558299e-12, -1.3717421e-14, -3.8400000e-16])
 
     def test_lennard_jones_sigma_epsilon_energy(self):
-        a = forcefields.lennard_jones_sigma_epsilon(2.0, [1.0, 0.25])
-        assert_almost_equal(a.energy, -0.015380859)
-        b = forcefields.lennard_jones_sigma_epsilon([2.0, 1.0], [1.0, 0.25])
-        assert_almost_equal(b.energy, [-0.015380859, 0])
-        c = forcefields.lennard_jones_sigma_epsilon(
-            [2.0, 1.0, 1.5], [0.5, 0.75])
-        assert_almost_equal(c.energy, [-0.0007322, -0.0461425, -0.0041095])
-        d = forcefields.lennard_jones_sigma_epsilon(
-            [400.0, 500.0, 600.0], [5e-10, 9e-9])
-        assert_almost_equal(d.energy, [0, 0, 0])
+        a = forcefields.lennard_jones_sigma_epsilon([1.0, 0.25])
+        assert_almost_equal(a.energy(2.0), -0.015380859)
+        b = forcefields.lennard_jones_sigma_epsilon([1.0, 0.25])
+        assert_almost_equal(b.energy([2.0, 1.0]), [-0.015380859, 0])
+        c = forcefields.lennard_jones_sigma_epsilon([0.5, 0.75])
+        assert_almost_equal(c.energy([2.0, 1.0, 1.5]), [-0.0007322, -0.0461425, -0.0041095])
+        d = forcefields.lennard_jones_sigma_epsilon([5e-10, 9e-9])
+        assert_almost_equal(d.energy([400.0, 500.0, 600.0]), [0, 0, 0])
 
     def test_lennard_jones_sigma_epsilon_force(self):
-        a = forcefields.lennard_jones_sigma_epsilon(2.0, [1.0, 0.25])
-        assert_almost_equal(a.force, -0.045410156)
-        b = forcefields.lennard_jones_sigma_epsilon([2.0, 4.0], [1.0, 0.25])
-        assert_almost_equal(b.force, [-0.0454102, -0.000366])
-        c = forcefields.lennard_jones_sigma_epsilon([3.0, 4.0], [3.0, 1.0])
-        assert_almost_equal(c.force, [8.0, -0.6877549])
+        a = forcefields.lennard_jones_sigma_epsilon([1.0, 0.25])
+        assert_almost_equal(a.force(2.0), -0.045410156)
+        b = forcefields.lennard_jones_sigma_epsilon([1.0, 0.25])
+        assert_almost_equal(b.force([2.0, 4.0]), [-0.0454102, -0.000366])
+        c = forcefields.lennard_jones_sigma_epsilon([3.0, 1.0])
+        assert_almost_equal(c.force([3.0, 4.0]), [8.0, -0.6877549])
 
     def test_buckingham_energy(self):
-        a = forcefields.buckingham(2.0, [1.0, 1.0, 1.0])
-        assert_almost_equal(a.energy, 0.1197103832)
-        b = forcefields.buckingham([2.0], [1.0, 1.0, 1.0])
-        assert_almost_equal(b.energy, 0.1197103832)
-        c = forcefields.buckingham([2.0, 4.0], [1.0, 1.0, 1.0])
-        assert_almost_equal(c.energy, [0.1197103832, 0.0180715])
-        d = forcefields.buckingham([2.0, 4.0, 5.0], [0.01, 0.01, 0.01])
-        assert_almost_equal(d.energy, [0.0096457, 0.0096055, 0.0095117])
+        a = forcefields.buckingham([1.0, 1.0, 1.0])
+        assert_almost_equal(a.energy(2.0), 0.1197103832)
+        b = forcefields.buckingham([1.0, 1.0, 1.0])
+        assert_almost_equal(b.energy([2.0]), 0.1197103832)
+        c = forcefields.buckingham([1.0, 1.0, 1.0])
+        assert_almost_equal(c.energy([2.0, 4.0]), [0.1197103832, 0.0180715])
+        d = forcefields.buckingham([0.01, 0.01, 0.01])
+        assert_almost_equal(d.energy([2.0, 4.0, 5.0]), [0.0096457, 0.0096055, 0.0095117])
 
     def test_buckingham_force(self):
-        a = forcefields.buckingham(2.0, [1.0, 1.0, 1.0])
-        assert_almost_equal(a.force, 0.08846028324)
-        b = forcefields.buckingham([2.0], [1.0, 1.0, 1.0])
-        assert_almost_equal(b.force, 0.08846028324)
-        c = forcefields.buckingham([2.0, 1.0, 4.0], [1.5, 0.1, 2.0])
-        assert_almost_equal(c.force, [0.0290596, -11.8642744, 0.0998156])
+        a = forcefields.buckingham([1.0, 1.0, 1.0])
+        assert_almost_equal(a.force(2.0), 0.08846028324)
+        b = forcefields.buckingham([1.0, 1.0, 1.0])
+        assert_almost_equal(b.force([2.0]), 0.08846028324)
+        c = forcefields.buckingham([1.5, 0.1, 2.0])
+        assert_almost_equal(c.force([2.0, 1.0, 4.0]), [0.0290596, -11.8642744, 0.0998156])
 
     def test_square_well_energy(self):
-        a = forcefields.square_well(2.0, [1.0, 1.5, 2.0])
-        assert_equal(a.energy, -1.0)
-        b = forcefields.square_well(0.5, [1.0, 2.0, 1.25])
-        assert_equal(b.energy, float('inf'))
-        c = forcefields.square_well(3.0, [0.5, 1.5, 1.25])
-        assert_equal(c.energy, 0)
-        d = forcefields.square_well([2.0, 0.5], [1.0, 1.5, 2.0])
-        assert_equal(d.energy, [-1.0, float('inf')])
-        e = forcefields.square_well([3.0, 3.0, 0.25], [1.0, 1.5, 1.25])
-        assert_equal(e.energy, [0, 0, float('inf')])
-        f = forcefields.square_well([3.0, 3.0, 0.25], [1.0, 1.5, 1.25], max_val=5000)
-        assert_equal(f.energy, [0, 0, 5000])
+        a = forcefields.square_well([1.0, 1.5, 2.0])
+        assert_equal(a.energy(2.0), -1.0)
+        b = forcefields.square_well([1.0, 2.0, 1.25])
+        assert_equal(b.energy(0.5), float('inf'))
+        c = forcefields.square_well([0.5, 1.5, 1.25])
+        assert_equal(c.energy(3.0), 0)
+        d = forcefields.square_well([1.0, 1.5, 2.0])
+        assert_equal(d.energy([2.0, 0.5]), [-1.0, float('inf')])
+        e = forcefields.square_well([1.0, 1.5, 1.25])
+        assert_equal(e.energy([3.0, 3.0, 0.25]), [0, 0, float('inf')])
+        f = forcefields.square_well([1.0, 1.5, 1.25], max_val=5000)
+        assert_equal(f.energy([3.0, 3.0, 0.25]), [0, 0, 5000])
 
     def test_square_well_force(self):
-        a = forcefields.square_well(2.0, [1.0, 1.5, 2.0])
+        a = forcefields.square_well([1.0, 1.5, 2.0])
         with self.assertRaises(ValueError):
-            a.update_force()
-        b = forcefields.square_well([2.0], [1.0, 1.5, 2.0])
+            a.force()
+        b = forcefields.square_well([1.0, 1.5, 2.0])
         with self.assertRaises(ValueError):
-            b.update_force()
+            b.force()
 
 
 if __name__ == '__main__':