tweaks to linear assignment

Former-commit-id: f880479eb34648a056b6d5a95126a59924bf08fd [formerly 5ad2525]
Former-commit-id: a04dfb80adc1b92fd013123e4c6ce0686ad793c6

pymatgen/optimization/linear_assignment.py

@@ -16,6 +16,7 @@
 
 import numpy as np
 
+from six.moves import range
 
 class LinearAssignment(object):
     """
@@ -120,8 +121,12 @@ def _augmenting_row_reduction(self):
         """
         unassigned = np.where(self._x == -1)[0]
         for i in unassigned:
-            while True:
-                #find smallest 2 values and indices
+            for _ in range(self.c.size):
+                # Time in this loop can be proportional to 1/epsilon
+                # This step is not strictly necessary, so cutoff early
+                # to avoid near-infinite loops
+
+                # find smallest 2 values and indices
                 temp = self.c[i] - self._v
                 j1 = np.argmin(temp)
                 u1 = temp[j1]
@@ -226,7 +231,7 @@ def _build_tree(self):
             #update predecessors
             _pred[shorter] = i
 
-            for j in np.argwhere(np.logical_and(self._d == mu, _todo)).flatten():
+            for j in np.nonzero(np.logical_and(self._d == mu, _todo))[0]:
                 if self._y[j] == -1:
                     return _pred, _ready, istar, j, mu
                 _scan[j] = True