fixing tests

arm61 · Jul 28, 2024 · fe87f9f · fe87f9f
1 parent 974f23f
commit fe87f9f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 9 deletions.
diff --git 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
@@ -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(