diff --git a/pylj/forcefields.py b/pylj/forcefields.py index bc6743e..4c5db74 100644 --- a/pylj/forcefields.py +++ b/pylj/forcefields.py @@ -154,8 +154,10 @@ def energy(self, dr): """ energy = self.a * np.exp(- np.multiply(self.b, dr)) - self.c / np.power(dr, 6) # Cut out infinite values where r = 0 - energy[energy > 10e300] = 0 - energy[energy < -10e300] = 0 + if type(dr) != float: + energy = np.array(energy) + energy[np.where(energy > 10e300)] = 0 + energy[np.where(energy < -10e300)] = 0 self.energy = energy return self.energy @@ -177,8 +179,10 @@ def force(self, dr): """ force = self.a * self.b * np.exp(- np.multiply(self.b, dr)) - 6 * self.c / np.power(dr, 7) # Cut out infinite values where r = 0 - force[force > 10e300] = 0 - force[force < -10e300] = 0 + if type(dr) != float: + force = np.array(force) + force[np.where(force > 10e300)] = 0 + force[np.where(force < -10e300)] = 0 self.force = force return self.force diff --git a/pylj/tests/test_util.py b/pylj/tests/test_util.py index df8c474..4c7c849 100644 --- a/pylj/tests/test_util.py +++ b/pylj/tests/test_util.py @@ -11,7 +11,7 @@ def test_system_square(self): 300, 8, mass=39.948, - constants=[1.363e-134, 9.273e-78], + constants=[[1.363e-134, 9.273e-78]], forcefield=ff.lennard_jones, ) assert_equal(a.number_of_particles, 2) @@ -34,7 +34,7 @@ def test_system_random(self): 8, init_conf="random", mass=39.948, - constants=[1.363e-134, 9.273e-78], + constants=[[1.363e-134, 9.273e-78]], forcefield=ff.lennard_jones, ) assert_equal(a.number_of_particles, 2) @@ -57,7 +57,7 @@ def test_system_too_big(self): 300, 1000, mass=39.948, - constants=[1.363e-134, 9.273e-78], + constants=[[1.363e-134, 9.273e-78]], forcefield=ff.lennard_jones, ) self.assertTrue( @@ -73,7 +73,7 @@ def test_system_too_small(self): 300, 2, mass=39.948, - constants=[1.363e-134, 9.273e-78], + constants=[[1.363e-134, 9.273e-78]], forcefield=ff.lennard_jones, ) self.assertTrue( @@ -89,7 +89,7 @@ def test_system_init_conf(self): 100, init_conf="horseradish", mass=39.948, - constants=[1.363e-134, 9.273e-78], + constants=[[1.363e-134, 9.273e-78]], forcefield=ff.lennard_jones, ) self.assertTrue(