Optimization: Replace loop with bvec[:, indices] for slicing rows

[SaFE-APIOpt]

AtlasToDiffusionReg/src/CODE/extract_single_shell.py

Line 101 in cfb76fb

new_dir = np.take(x, indices)

Current code:

dirs = []
for x in bvec:
    new_dir = np.take(x, indices)
    dirs.append(new_dir)
new_bvec = np.stack(dirs, axis=0)

Suggested replacement:
new_bvec = bvec[:, indices]
The current implementation loops over all rows in bvec, calling np.take on each one, and appends results to a list before stacking. This introduces significant overhead from:

Python-level iteration;

Repeated function calls;

Intermediate list and array construction.

The optimized version uses NumPy’s advanced indexing to extract the desired columns from all rows at once:

It is fully vectorized, much faster, and more memory efficient;

The result is exactly the same.

[larsenv]

good idea