Final black formatting and recursion fix for PKDTree

Author: Aryaman-Chaudhri

package/MDAnalysis/lib/pkdtree.py

@@ -56,9 +56,7 @@ class AugmentedPKDTree(object):
     duplicate particles while maintaining correctness.
     """
 
-    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
@@ -78,15 +76,15 @@ def __init__(
         self.leafsize = leafsize
         self.dim = 3  # 3D systems
         self.box = box
-        
+
         if self.box is not None:
             self.box = np.asarray(box, dtype=np.float32)
             # Cache orthorhombic check - don't recompute every time
             self._is_ortho = self._check_orthorhombic()
         else:
             self.box = None
             self._is_ortho = True
-            
+
         self._built = False
         self.cutoff: Optional[float] = None
         self.mapping = None  # For compatibility, always None in this implementation
@@ -113,9 +111,7 @@ 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
 
         For orthorhombic boxes, uses scipy's native boxsize parameter
@@ -139,29 +135,33 @@ def set_coords(
         # Set coords dtype to float32
         coords = np.asarray(coords, dtype=np.float32)
         self.cutoff = cutoff
-        
+
         if self.box is None:
             # Non-periodic case
             if cutoff is not None:
-                raise RuntimeError("Donot provide cutoff distance for non PBC aware calculations")
-            
+                raise RuntimeError(
+                    "Donot provide cutoff distance for non PBC aware calculations"
+                )
+
             self.coords = coords
             self.ckdt = cKDTree(self.coords, leafsize=self.leafsize)
             self._use_scipy = True
             self._use_capped = False
             self._built = True
             return
-        
+
         # Periodic case - cutoff required
         if cutoff is None:
             raise RuntimeError("Provide a cutoff distance with tree.set_coords(...)")
-            
+
         if self._is_ortho:
             # Orthorhombic box - use scipy's native PBC support
             # Wrap coordinates once during tree construction
             # scipy's boxsize parameter handles PBC during queries efficiently
             self.coords = apply_PBC(coords, self.box)
-            self.ckdt = cKDTree(self.coords, leafsize=self.leafsize, boxsize=self.box[:3])
+            self.ckdt = cKDTree(
+                self.coords, leafsize=self.leafsize, boxsize=self.box[:3]
+            )
             self._use_scipy = True
             self._use_capped = False
         else:
@@ -197,7 +197,7 @@ def search(self, centers: npt.ArrayLike, radius: float) -> npt.NDArray:
                 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.")
-            
+
             if self._use_scipy:
                 # scipy's boxsize parameter handles PBC, but wrap centers to be safe
                 wrapped_centers = apply_PBC(centers, self.box)
@@ -215,18 +215,22 @@ def search(self, centers: npt.ArrayLike, radius: float) -> npt.NDArray:
                         self.coords,
                         radius,
                         box=self.box,
-                        return_distances=False
+                        return_distances=False,
                     )
                     if len(pairs) > 0:
                         all_indices.extend(pairs[:, 1])
-                self._indices = np.array(all_indices, dtype=np.intp) if all_indices else np.array([], dtype=np.intp)
+                self._indices = (
+                    np.array(all_indices, dtype=np.intp)
+                    if all_indices
+                    else np.array([], dtype=np.intp)
+                )
         else:
             # Non-periodic case
             indices = list(self.ckdt.query_ball_point(centers, radius))
             self._indices = np.array(
                 list(itertools.chain.from_iterable(indices)), dtype=np.intp
             )
-        
+
         if self._indices.size > 0:
             self._indices = np.asarray(unique_int_1d(self._indices))
         return self._indices
@@ -262,13 +266,15 @@ def search_pairs(self, radius: float) -> npt.NDArray:
                 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.")
-        
+
         if self._use_scipy:
             # scipy handles PBC with boxsize parameter
             pairs = np.array(list(self.ckdt.query_pairs(radius)), dtype=np.intp)
         elif self._use_capped:
             # Use self_capped_distance for triclinic boxes
-            pairs = self_capped_distance(self.coords, radius, box=self.box, return_distances=False)
+            pairs = self_capped_distance(
+                self.coords, radius, box=self.box, return_distances=False
+            )
             pairs = np.asarray(pairs)
         else:
             # Non-periodic fallback
@@ -319,40 +325,49 @@ def search_tree(self, centers: npt.ArrayLike, radius: float) -> np.ndarray:
                 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.")
-            
+
         if self._use_scipy:
             if self.pbc:
                 # Wrap centers for scipy's boxsize parameter
                 wrapped_centers = apply_PBC(centers, self.box)
-                other_tree = cKDTree(wrapped_centers, leafsize=self.leafsize, boxsize=self.box[:3])
+                other_tree = cKDTree(
+                    wrapped_centers, leafsize=self.leafsize, boxsize=self.box[:3]
+                )
             else:
                 other_tree = cKDTree(centers, leafsize=self.leafsize)
-            
+
             pairs_list = other_tree.query_ball_tree(self.ckdt, radius)
-            pairs = np.array([[i, j] for i, lst in enumerate(pairs_list) for j in lst], dtype=np.intp)
-            
+            pairs = np.array(
+                [[i, j] for i, lst in enumerate(pairs_list) for j in lst], dtype=np.intp
+            )
+
         elif self._use_capped:
             # Wrap centers for triclinic boxes and use capped_distance
             wrapped_centers = apply_PBC(centers, self.box)
-            pairs = capped_distance(wrapped_centers, self.coords, radius, 
-                                   box=self.box, return_distances=False)
+            pairs = capped_distance(
+                wrapped_centers,
+                self.coords,
+                radius,
+                box=self.box,
+                return_distances=False,
+            )
             pairs = np.asarray(pairs)
         else:
             # Non-periodic fallback
             other_tree = cKDTree(centers, leafsize=self.leafsize)
             pairs_list = other_tree.query_ball_tree(self.ckdt, radius)
-            pairs = np.array([[i, j] for i, lst in enumerate(pairs_list) for j in lst], dtype=np.intp)
-        
+            pairs = np.array(
+                [[i, j] for i, lst in enumerate(pairs_list) for j in lst], dtype=np.intp
+            )
+
         if pairs.size > 0:
             pairs = unique_rows(pairs)
         return pairs
-    
+
 
 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:
         """
 
         Parameters
@@ -373,9 +388,9 @@ def __init__(
         self.dim = 3  # 3D systems
         self.box = box
         self._built = False
-        self.cutoff: Optional[float] = None 
+        self.cutoff: Optional[float] = None
         self.mapping = None  # For compatibility, always None in this implementation
-        
+
         _use_augmented = False
         if box is not None:
             box_array = np.asarray(box, dtype=np.float32)
@@ -384,11 +399,13 @@ def __init__(
                     _use_augmented = True
             else:
                 _use_augmented = True
-        
+
         self._use_augmented = _use_augmented
-        
+
         if self._use_augmented:
-            self._tree: Union[AugmentedPKDTree, cKDTree, None] = AugmentedPKDTree(box=self.box, leafsize=leafsize)
+            self._tree: Union[AugmentedPKDTree, cKDTree, None] = AugmentedPKDTree(
+                box=self.box, leafsize=leafsize
+            )
         else:
             self._tree = None
             if box is not None:
@@ -406,9 +423,7 @@ 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
@@ -427,19 +442,25 @@ def set_coords(
         else:
             coords = np.asarray(coords, dtype=np.float32)
             self.cutoff = cutoff
-            
+
             if self.box is None:
                 if cutoff is not None:
-                    raise RuntimeError("Donot provide cutoff distance for non PBC aware calculations")
+                    raise RuntimeError(
+                        "Donot provide cutoff distance for non PBC aware calculations"
+                    )
                 self.coords = coords
                 self._tree = cKDTree(self.coords, leafsize=self.leafsize)
             else:
                 if cutoff is None:
-                    raise RuntimeError("Provide a cutoff distance with tree.set_coords(...)")
+                    raise RuntimeError(
+                        "Provide a cutoff distance with tree.set_coords(...)"
+                    )
                 self.coords = apply_PBC(coords, self.box)
                 box_array = np.asarray(self.box, dtype=np.float32)
-                self._tree = cKDTree(self.coords, leafsize=self.leafsize, boxsize=box_array[:3])
-            
+                self._tree = cKDTree(
+                    self.coords, leafsize=self.leafsize, boxsize=box_array[:3]
+                )
+
             self._built = True
 
     def search(self, centers: npt.ArrayLike, radius: float) -> npt.NDArray:
@@ -458,7 +479,7 @@ def search(self, centers: npt.ArrayLike, radius: float) -> npt.NDArray:
         if self._use_augmented:
             assert self._tree is not None
             return self._tree.search(centers, radius)
-        
+
         centers = np.asarray(centers, dtype=np.float32)
         if centers.shape == (self.dim,):
             centers = centers.reshape((1, self.dim))
@@ -474,11 +495,11 @@ def search(self, centers: npt.ArrayLike, radius: float) -> npt.NDArray:
         else:
             assert isinstance(self._tree, cKDTree)
             indices = list(self._tree.query_ball_point(centers, radius))
-        
+
         self._indices = np.array(
             list(itertools.chain.from_iterable(indices)), dtype=np.intp
         )
-        
+
         if self._indices.size > 0:
             self._indices = np.asarray(unique_int_1d(self._indices))
         return self._indices
@@ -515,16 +536,16 @@ def search_pairs(self, radius: float) -> npt.NDArray:
         if self._use_augmented:
             assert self._tree is not None
             return self._tree.search_pairs(radius)
-        
+
         if self.pbc:
             if self.cutoff is None:
                 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.")
-        
+
         assert isinstance(self._tree, cKDTree)
         pairs = np.array(list(self._tree.query_pairs(radius)), dtype=np.intp)
-        
+
         if pairs.size > 0:
             pairs = np.sort(pairs, axis=1)
             pairs = unique_rows(pairs)
@@ -556,7 +577,7 @@ def search_tree(self, centers: npt.ArrayLike, radius: float) -> np.ndarray:
         if self._use_augmented:
             assert self._tree is not None
             return self._tree.search_tree(centers, radius)
-    
+
         centers = np.asarray(centers, dtype=np.float32)
         if centers.shape == (self.dim,):
             centers = centers.reshape((1, self.dim))
@@ -568,13 +589,17 @@ def search_tree(self, centers: npt.ArrayLike, radius: float) -> np.ndarray:
                 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(wrapped_centers, leafsize=self.leafsize, boxsize=box_array[:3])
+            other_tree = cKDTree(
+                wrapped_centers, leafsize=self.leafsize, boxsize=box_array[:3]
+            )
         else:
             other_tree = cKDTree(centers, leafsize=self.leafsize)
-        
+
         pairs_list = other_tree.query_ball_tree(self._tree, radius)
-        pairs = np.array([[i, j] for i, lst in enumerate(pairs_list) for j in lst], dtype=np.intp)
+        pairs = np.array(
+            [[i, j] for i, lst in enumerate(pairs_list) for j in lst], dtype=np.intp
+        )
 
         if pairs.size > 0:
             pairs = unique_rows(pairs)
-        return pairs
+        return pairs

testsuite/MDAnalysisTests/lib/test_pkdtree.py

@@ -29,11 +29,13 @@
 from MDAnalysis.lib.pkdtree import PeriodicKDTree, AugmentedPKDTree
 from MDAnalysis.lib.distances import transform_StoR
 
+
 @pytest.fixture(params=[PeriodicKDTree, AugmentedPKDTree])
 def tree_class(request):
     """Fixture to run tests on both PeriodicKDTree and AugmentedPKDTree."""
     return request.param
 
+
 # fractional coordinates for data points
 f_dataset = np.array(
     [
@@ -62,7 +64,7 @@ def tree_class(request):
         ),
     ),
 )
-def test_setcoords(tree_class , b, cut, result):
+def test_setcoords(tree_class, b, cut, result):
     coords = np.array([[1, 1, 1], [2, 2, 2]], dtype=np.float32)
     if b is not None:
         b = np.array(b, dtype=np.float32)
@@ -93,7 +95,7 @@ def test_searchfail(tree_class):
         ([10, 10, 10, 45, 60, 90], [2.1, -3.1, 0.1], [2, 3]),
     ),
 )
-def test_search(tree_class , b, q, result):
+def test_search(tree_class, b, q, result):
     b = np.array(b, dtype=np.float32)
     q = transform_StoR(np.array(q, dtype=np.float32), b)
     cutoff = 3.0
@@ -165,7 +167,7 @@ def test_ckd_searchpairs_nopbc(tree_class, radius, result):
                                                                        [0, 3]])
                          ))
 # fmt: on
-def test_searchtree(tree_class , b, q, result):
+def test_searchtree(tree_class, b, q, result):
     b = np.array(b, dtype=np.float32)
     cutoff = 3.0
     coords = transform_StoR(f_dataset, b)