Commit 7f1c7ed2 authored by Michiel Cottaar's avatar Michiel Cottaar
Browse files

ENH: more efficiently compute idx_centroid for precomputed grid

parent 2407d7c4
Pipeline #5294 failed with stage
in 9 minutes and 54 seconds
......@@ -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:])
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment