Hey,
I have this procedure, where I place a number 3D volumes within a larger Volume at different positions.
Due to the for-loop this is currently a bottleneck on speed. I feel like there has to be a way to do vectorize this but I can’t quite get it to work.
(I took out any consideration of edge cases to make the example more concise).
# Tensor containing N different 3D volumes. We want to place these at certain spots in a bigger volume.
N, psf_h, psf_w, psf_d = sub_vols.shape
# b, ch, x, y, z are lists that specify where each sub volume should be placed.
for idx in range(N):
big_volume[b[idx], ch[idx],
z[idx] : z[idx]+psf_h,
y[idx] : y[idx]+psf_w,
x[idx] : x[idx]+psf_d] += sub_vols[idx]