I'm have a 3D problem where to final output is an array in the xy plane. I have an array in the x-z plane (dimensions (xsiz, zsiz)) and an array in the y-plane (dimension ysiz) as below:

xz = np.zeros((xsiz, zsiz))
y = (np.arange(ysiz)*(zsiz/ysiz)).astype(int)

xz can be thought of as an array of (zsiz) column vectors of size (xsiz) and labelled by z in range (0, zsiz-1). These are not conveniently accessible given the current setup - I've been retrieving them by np.transpose(xz)[z]. I would like the y array to act like a list of z values and take the column vectors labelled by these z values and combine them in a matrix with final dimension (xsiz, ysiz). (It seems likely to me that it will be easier to work with the transpose of xz so the row vectors can be retrieved as above and combined giving a (ysiz, xsiz) matrix which can then be transposed but I may be wrong.)

This would be a simple using for loops and I've given an example of a such a loop that does what I want below in case my explanation isn't clear. However, the final intention is for this code to be parallelized using CuPy so ideally I would like the entire process to be carried out by matrix manipulation. It seems like it should be possible like this but I can't think how!

Any help greatly appreciated.

import numpy as np

xsiz = 5 #sizes given random values for example
ysiz = 6
zsiz = 4


xz = np.arange(xsiz*zsiz).reshape(xsiz, zsiz)
y = (np.arange(ysiz)*(zsiz/ysiz)).astype(int)

xzT = np.transpose(xz)

final_xyT = np.zeros((0, xsiz))

for i in range(ysiz):
    index = y[i]
    xvec = xzT[index]
    final_xyT = np.vstack((final_xyT, xvec))

#indexing could go wrong here if y contained large numbers
#CuPy's indexing wraps around so hopefully this shouldn't be too big an issue

final_xy = np.transpose(final_xyT)
print(xz)
print(final_xy)

CodePudding user response：

If I correctly get your problem you need this:

xz[:,y]