shenfun
shenfun copied to clipboard
data layout
Hi,
I'm trying to understand the layout of data among MPI procs. Could anybody please take a look at the following pics in order to see if I got it right?
Hi
The array layouts are actually possible to configure when creating the TensorProductSpace
. The keyword argument axes
can be used to determine the order of which the different axes are transformed, and this determines also the decomposition. Furthermore, slab=True
can be used to decompose in a 1- instead of 2-dimensional process grid. The default is to use the pencil method, though, which is what you are drawing.
From what I can tell the 2D pics are correct with default settings. I'm also pretty sure the 3D plots are correct. In case you're not aware, you can double check by printing the parts of the global array that belongs on each cpu as
print(comm.Get_rank(), T.local_slice(spectral=False))
print(comm.Get_rank(), T.local_slice(spectral=True))
Ok, thanks, I also found local_wavenumbers
and local_mesh
to be useful in order to inspect the layout:
In 2D:
X = T.local_mesh(True)
K = T.local_wavenumbers()
print(comm.Get_rank(), ' spectral kx=', K[0][[0, -1]], ', ky =', K[1][0, [0, -1]])
print(comm.Get_rank(), ' physical x=', X[0][[0, -1], 0], ', y =', X[1][0, [0, -1]])
And in 3D:
X = T.local_mesh(True)
K = T.local_wavenumbers()
print(comm.Get_rank(), ' spectral kx=',K[0][[0,-1],0,0],', ky =',K[1][0,[0,-1],0],' kz =',K[2][0,0,[0,-1]])
print(comm.Get_rank(), ' physical x=',X[0][[0,-1],0,0],', y =',X[1][0,[0,-1],0],' z =',X[2][0,0,[0,-1]])