ENH: more efficiently compute idx_centroid for precomputed grid

gyral_structure/basis/radial.py

@@ -159,13 +159,14 @@ class RadialBasis(BasisFunc):
             use = (voxels >= 0).all(-1) & (voxels < intersects.shape).all(-1)
             list_of_lists[use] = intersects[tuple(voxels[use].T)]
             list_of_lists = list(list_of_lists)
+            idx_centroid = sp.concatenate(list_of_lists)
         else:
             tree2 = spatial.cKDTree(req.center())
             size1 = sp.asarray(self.size)
             max_size = size1.max() + req.radius().max()
             list_of_lists = tree2.query_ball_tree(self.tree, r=max_size)
+            idx_centroid = sp.fromiter(itertools.chain(*list_of_lists), dtype='i4')
         idx_req_compressed = sp.append(0, sp.cumsum([len(l) for l in list_of_lists]))
-        idx_centroid = sp.fromiter(itertools.chain(*list_of_lists), dtype='i4')
         if return_compressed:
             return idx_req_compressed, idx_centroid
         idx_req = sp.digitize(sp.arange(idx_req_compressed[-1]), idx_req_compressed[1:])
@@ -205,7 +206,7 @@ class RadialBasis(BasisFunc):
         center_tree = spatial.cKDTree(coords_center.reshape(3, -1).T)
         intersects = np.zeros(np.prod(coords_center.shape[1:]), dtype='object')
         for idx, neighbours in enumerate(center_tree.query_ball_tree(self.tree, r=max_size)):
-            intersects[idx] = neighbours
+            intersects[idx] = np.array(neighbours, dtype='i4')
 
         self._precomputed_grids = (
             max_size, sp.linalg.inv(affine), intersects.reshape(coords_center.shape[1:])