bug fix: NaNs in pairwise forces

johnaparker · Feb 5, 2020 · b2d7f66 · b2d7f66
1 parent f11c3c2
commit b2d7f66
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/stoked/forces.py b/stoked/forces.py
@@ -17,7 +17,7 @@ def force(self):
 
         r_ij = self.position[:,np.newaxis] - self.position[np.newaxis]   # N x N x 3
         F_ij = self.force_func(r_ij)   # N x N x 3
-        np.einsum('iix->x', F_ij)[...] = 0
+        np.einsum('iix->ix', F_ij)[...] = 0
         F_i = np.sum(F_ij, axis=1)
 
         return F_i