fixing linters continued

Author: Aryaman-Chaudhri

benchmarks/benchmarks/neighbors.py

@@ -12,10 +12,12 @@ class NeighborsBench:
 
     def setup(self, number_of_atoms, cutoff):
         """Setup called before each benchmark with each parameter combination."""
-        self.box = np.array([170.0, 70.0, 120.0, 90.0, 90.0, 90.0], dtype=np.float32)
-        self.positions = (np.random.rand(number_of_atoms, 3) * self.box[:3]).astype(
-            np.float32
+        self.box = np.array(
+            [170.0, 70.0, 120.0, 90.0, 90.0, 90.0], dtype=np.float32
         )
+        self.positions = (
+            np.random.rand(number_of_atoms, 3) * self.box[:3]
+        ).astype(np.float32)
         self.centre = (self.box[:3] / 2.0).reshape(1, 3)
         self.cutoff = cutoff
 

package/MDAnalysis/lib/pkdtree.py

@@ -62,7 +62,9 @@ class AugmentedPKDTree(object):
 
     """
 
-    def __init__(self, box: Optional[npt.ArrayLike] = None, leafsize: int = 10) -> None:
+    def __init__(
+        self, box: Optional[npt.ArrayLike] = None, leafsize: int = 10
+    ) -> None:
         """
 
         Parameters
@@ -96,7 +98,9 @@ def pbc(self):
         """
         return self.box is not None
 
-    def set_coords(self, coords: npt.ArrayLike, cutoff: Optional[float] = None) -> None:
+    def set_coords(
+        self, coords: npt.ArrayLike, cutoff: Optional[float] = None
+    ) -> None:
         """Constructs KDTree from the coordinates
 
         Wrapping of coordinates to the primary unit cell is enforced
@@ -143,15 +147,18 @@ def set_coords(self, coords: npt.ArrayLike, cutoff: Optional[float] = None) -> N
             # Bring the coordinates in the central cell
             self.coords = apply_PBC(coords, self.box)
             # generate duplicate images
-            self.aug, self.mapping = augment_coordinates(self.coords, self.box, cutoff)
+            self.aug, self.mapping = augment_coordinates(
+                self.coords, self.box, cutoff
+            )
             # Images + coords
             self.all_coords = np.concatenate([self.coords, self.aug])
             self.ckdt = cKDTree(self.all_coords, leafsize=self.leafsize)
         else:
             # if cutoff distance is provided for non PBC calculations
             if cutoff is not None:
                 raise RuntimeError(
-                    "Donot provide cutoff distance for" " non PBC aware calculations"
+                    "Donot provide cutoff distance for"
+                    " non PBC aware calculations"
                 )
             self.coords = coords
             self.ckdt = cKDTree(self.coords, self.leafsize)
@@ -182,9 +189,13 @@ def search(self, centers: npt.ArrayLike, radius: float) -> npt.NDArray:
         # Sanity check
         if self.pbc:
             if self.cutoff is None:
-                raise ValueError("Cutoff needs to be provided when working with PBC.")
+                raise ValueError(
+                    "Cutoff needs to be provided when working with PBC."
+                )
             if self.cutoff < radius:
-                raise RuntimeError("Set cutoff greater or equal to the radius.")
+                raise RuntimeError(
+                    "Set cutoff greater or equal to the radius."
+                )
             # Bring all query points to the central cell
             wrapped_centers = apply_PBC(centers, self.box)
             indices = list(self.ckdt.query_ball_point(wrapped_centers, radius))
@@ -232,15 +243,23 @@ def search_pairs(self, radius: float) -> npt.NDArray:
 
         if self.pbc:
             if self.cutoff is None:
-                raise ValueError("Cutoff needs to be provided when working with PBC.")
+                raise ValueError(
+                    "Cutoff needs to be provided when working with PBC."
+                )
             if self.cutoff < radius:
-                raise RuntimeError("Set cutoff greater or equal to the radius.")
+                raise RuntimeError(
+                    "Set cutoff greater or equal to the radius."
+                )
 
         pairs = np.array(list(self.ckdt.query_pairs(radius)), dtype=np.intp)
         if self.pbc:
             if len(pairs) > 1:
-                pairs[:, 0] = undo_augment(pairs[:, 0], self.mapping, len(self.coords))
-                pairs[:, 1] = undo_augment(pairs[:, 1], self.mapping, len(self.coords))
+                pairs[:, 0] = undo_augment(
+                    pairs[:, 0], self.mapping, len(self.coords)
+                )
+                pairs[:, 1] = undo_augment(
+                    pairs[:, 1], self.mapping, len(self.coords)
+                )
         if pairs.size > 0:
             # First sort the pairs then pick the unique pairs
             pairs = np.sort(pairs, axis=1)
@@ -288,9 +307,13 @@ class initialization
         # Sanity check
         if self.pbc:
             if self.cutoff is None:
-                raise ValueError("Cutoff needs to be provided when working with PBC.")
+                raise ValueError(
+                    "Cutoff needs to be provided when working with PBC."
+                )
             if self.cutoff < radius:
-                raise RuntimeError("Set cutoff greater or equal to the radius.")
+                raise RuntimeError(
+                    "Set cutoff greater or equal to the radius."
+                )
             # Bring all query points to the central cell
             wrapped_centers = apply_PBC(centers, self.box)
             other_tree = cKDTree(wrapped_centers, leafsize=self.leafsize)
@@ -300,7 +323,9 @@ class initialization
                 dtype=np.intp,
             )
             if pairs.size > 0:
-                pairs[:, 1] = undo_augment(pairs[:, 1], self.mapping, len(self.coords))
+                pairs[:, 1] = undo_augment(
+                    pairs[:, 1], self.mapping, len(self.coords)
+                )
         else:
             other_tree = cKDTree(centers, leafsize=self.leafsize)
             pairs = other_tree.query_ball_tree(self.ckdt, radius)
@@ -315,13 +340,17 @@ class initialization
 
 class PeriodicKDTree(object):
 
-    def __init__(self, box: Optional[npt.ArrayLike] = None, leafsize: int = 10) -> None:
+    def __init__(
+        self, box: Optional[npt.ArrayLike] = None, leafsize: int = 10
+    ) -> None:
         self.leafsize = leafsize
         self.dim = 3
         self.box = box
         self._built = False
-        self.cutoff = None
-        self.mapping = None
+
+        self.cutoff: Optional[float] = None
+        self.mapping: Optional[npt.NDArray] = None
+        self._tree: Optional[Union[AugmentedPKDTree, cKDTree]] = None
 
         _use_augmented = False
         if box is not None:
@@ -353,7 +382,9 @@ def pbc(self):
         """
         return self.box is not None
 
-    def set_coords(self, coords: npt.ArrayLike, cutoff: Optional[float] = None) -> None:
+    def set_coords(
+        self, coords: npt.ArrayLike, cutoff: Optional[float] = None
+    ) -> None:
         """Constructs KDTree from the coordinates
 
         Parameters
@@ -416,12 +447,18 @@ def search(self, centers: npt.ArrayLike, radius: float) -> npt.NDArray:
 
         if self.pbc:
             if self.cutoff is None:
-                raise ValueError("Cutoff needs to be provided when working with PBC.")
+                raise ValueError(
+                    "Cutoff needs to be provided when working with PBC."
+                )
             if self.cutoff < radius:
-                raise RuntimeError("Set cutoff greater or equal to the radius.")
+                raise RuntimeError(
+                    "Set cutoff greater or equal to the radius."
+                )
             wrapped_centers = apply_PBC(centers, self.box)
             assert isinstance(self._tree, cKDTree)
-            indices = list(self._tree.query_ball_point(wrapped_centers, radius))
+            indices = list(
+                self._tree.query_ball_point(wrapped_centers, radius)
+            )
         else:
             assert isinstance(self._tree, cKDTree)
             indices = list(self._tree.query_ball_point(centers, radius))
@@ -442,9 +479,6 @@ def get_indices(self) -> npt.NDArray:
         indices : NDArray
           neighbors for the last query points and search radius
         """
-        if self._use_augmented:
-            assert self._tree is not None
-            return self._tree.get_indices()
         return self._indices
 
     def search_pairs(self, radius: float) -> npt.NDArray:
@@ -469,9 +503,13 @@ def search_pairs(self, radius: float) -> npt.NDArray:
 
         if self.pbc:
             if self.cutoff is None:
-                raise ValueError("Cutoff needs to be provided when working with PBC.")
+                raise ValueError(
+                    "Cutoff needs to be provided when working with PBC."
+                )
             if self.cutoff < radius:
-                raise RuntimeError("Set cutoff greater or equal to the radius.")
+                raise RuntimeError(
+                    "Set cutoff greater or equal to the radius."
+                )
 
         assert isinstance(self._tree, cKDTree)
         pairs = np.array(list(self._tree.query_pairs(radius)), dtype=np.intp)
@@ -514,9 +552,13 @@ def search_tree(self, centers: npt.ArrayLike, radius: float) -> np.ndarray:
 
         if self.pbc:
             if self.cutoff is None:
-                raise ValueError("Cutoff needs to be provided when working with PBC.")
+                raise ValueError(
+                    "Cutoff needs to be provided when working with PBC."
+                )
             if self.cutoff < radius:
-                raise RuntimeError("Set cutoff greater or equal to the radius.")
+                raise RuntimeError(
+                    "Set cutoff greater or equal to the radius."
+                )
             wrapped_centers = apply_PBC(centers, self.box)
             box_array = np.asarray(self.box, dtype=np.float32)
             other_tree = cKDTree(