updating forcefields method for improved particle mixing

pylj/forcefields.py

@@ -1,20 +1,35 @@
 import numpy as np
 
-
-class lennard_jones(object):
+class lennard_jones_base(object):
     r"""Calculate the energy or force for a pair of particles using the
-    Lennard-Jones (A/B variant) forcefield.
+    Lennard-Jones forcefield. Either the a_b or sigma_epsilon constants
+    must be specified
 
     Parameters
     ----------
-    constants: float, array_like
+    a_b: float, array_like
         An array of length two consisting of the A and B parameters for the
         12-6 Lennard-Jones function
-
+    sigma_epsilon: float, array_like
+        An array of length two consisting of the sigma and epsilon parameters for the
+        12-6 Lennard-Jones function
     """
-    def __init__(self, constants):
-        self.a = constants[0]
-        self.b = constants[1]
+    def __init__(self, constants, a_b = False, sigma_epsilon = False):
+
+        if sigma_epsilon:
+            self.sigma = constants[0]
+            self.epsilon = constants[1]
+            self.a = 4 * self.epsilon * (self.sigma**12)
+            self.b = 4 * self.epsilon * (self.sigma**6)
+            self.type = 'sigma_epsilon'
+
+        elif a_b:
+            self.a = constants[0]
+            self.b = constants[1]
+            self.sigma = (self.a / self.b)**(1/6)
+            self.epsilon = (self.b**2)/(4*self.a)
+            self.type = 'a_b'
+
 
     def energy(self, dr ):
         r"""Calculate the energy for a pair of particles using the
@@ -53,65 +68,49 @@ def force(self, dr):
         """
         self.force = 12 * self.a * np.power(dr, -13) - (6 * self.b * np.power(dr, -7))
         return self.force
+
+    def mixing(self, constants_2):
+
+        if self.type == 'a_b':
+            a2 = constants_2[0]
+            b2 = constants_2[1]
+            sigma2 = (a2 / b2)**(1/6)
+            epsilon2 = (b2**2)/(4*a2)
+
+        elif self.type == 'sigma_epsilon':
+            sigma2 = constants_2[0]
+            epsilon2 = constants_2[1]
+
+        self.sigma = (self.sigma+sigma2)/2
+        self.epsilon = np.sqrt(self.epsilon**2 + epsilon2**2)
+        self.a = 4 * self.epsilon * (self.sigma**12)
+        self.b = 4 * self.epsilon * (self.sigma**6)
 
+class lennard_jones(lennard_jones_base):
+    r"""Calculate the energy or force for a pair of particles using the
+    Lennard-Jones (A/B variant) forcefield. Maps to lennard_jones_base class
 
+    Parameters
+    ----------
+    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, constants):
+        super().__init__(constants, a_b = True)
 
-class lennard_jones_sigma_epsilon(object):
+class lennard_jones_sigma_epsilon(lennard_jones_base):
     r"""Calculate the energy or force for a pair of particles using the
-    Lennard-Jones (sigma/epsilon variant) forcefield.
+    Lennard-Jones (sigma/epsilon variant) forcefield. Maps to lennard_jones_base class
 
     Parameters
     ----------
     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, constants):
-        self.sigma = constants[0]
-        self.epsilon = constants[1]
-
-    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(dr, -12) - (
-                        4 * self.epsilon * np.power(self.sigma, 6) * np.power(dr, -6)) 
-        return self.energy 
-
-    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(
-            dr, -13) - (24 * self.epsilon * np.power(self.sigma, 6) * np.power(dr, -7))
-        return self.force
-
-
+        super().__init__(constants, sigma_epsilon = True)
 
 class buckingham(object):
     r""" Calculate the energy or force for a pair of particles using the

pylj/pairwise.py

@@ -52,13 +52,13 @@ def compute_force(particles, box_length, cut_off, constants, forcefield, mass):
         for i in range(len(distances)):
             if pair != pair_types[i]:
                 type_distances[i] = 0
-        if pair.split(',')[0] == pair.split(',')[1]:
-            constants_type = constants[int(pair.split(',')[0])]
-        else:
-            constants_1 = np.array(constants[int(pair.split(',')[0])])
+        type_1 = pair.split(',')[0]
+        type_2 = pair.split(',')[1]
+        constants_1 = np.array(constants[int(pair.split(',')[0])])
+        ff = forcefield(constants_1)
+        if type_1 != type_2:
             constants_2 = np.array(constants[int(pair.split(',')[1])])
-            constants_type = np.sqrt(constants_1*constants_2)
-        ff = forcefield(constants_type)
+            ff.mixing(constants_2)
         type_forces = ff.force(type_distances)
         type_energies = ff.energy(distances)
         type_forces = np.nan_to_num(type_forces)
@@ -233,14 +233,16 @@ def compute_energy(particles, box_length, cut_off, constants, forcefield):
         for i in range(len(distances)):
             if pair != pair_types[i]:
                 type_distances[i] = 0
-        if pair.split(',')[0] == pair.split(',')[1]:
-            constants_type = constants[int(pair.split(',')[0])]
-        else:
-            constants_1 = np.array(constants[int(pair.split(',')[0])])
+        type_1 = pair.split(',')[0]
+        type_2 = pair.split(',')[1]
+        constants_1 = np.array(constants[int(pair.split(',')[0])])
+        ff = forcefield(constants_1)
+        if type_1 != type_2:
             constants_2 = np.array(constants[int(pair.split(',')[1])])
-            constants_type = np.sqrt(constants_1*constants_2)
-        ff = forcefield(constants_type)
+            ff.mixing(constants_2)
+        type_forces = ff.force(type_distances)
         type_energies = ff.energy(distances)
+        type_forces = np.nan_to_num(type_forces)
         type_energies = np.nan_to_num(type_energies)
         energies+=type_energies
     energies[np.where(distances > cut_off)] = 0.0
@@ -287,15 +289,18 @@ def calculate_pressure(
         for i in range(len(distances)):
             if pair != pair_types[i]:
                 type_distances[i] = 0
-        if pair.split(',')[0] == pair.split(',')[1]:
-            constants_type = constants[int(pair.split(',')[0])]
-        else:
-            constants_1 = np.array(constants[int(pair.split(',')[0])])
+        type_1 = pair.split(',')[0]
+        type_2 = pair.split(',')[1]
+        constants_1 = np.array(constants[int(pair.split(',')[0])])
+        ff = forcefield(constants_1)
+        if type_1 != type_2:
             constants_2 = np.array(constants[int(pair.split(',')[1])])
-            constants_type = np.sqrt(constants_1*constants_2)
-        ff = forcefield(constants_type)
+            ff.mixing(constants_2)
+        type_forces = ff.force(type_distances)
         type_energies = ff.energy(distances)
+        type_forces = np.nan_to_num(type_forces)
         type_energies = np.nan_to_num(type_energies)
+        forces+=type_forces
         energies+=type_energies
     forces[np.where(distances > cut_off)] = 0.0
     pres = np.sum(forces * distances)

pylj/tests/test_pairwise.py

@@ -76,7 +76,7 @@ def test_calculate_pressure(self):
             constants=[[1.363e-134, 9.273e-78]],
             forcefield=ff.lennard_jones
         )
-        assert_almost_equal(p * 1e24, 9.20399999)
+        assert_almost_equal(p * 1e24, 7.07368867)
 
     def test_pbc_correction(self):
         a = pairwise.pbc_correction(1, 10)