def collect_orbitals(a_xo, coords, root=0):
"""Collect array distributed over orbitals to root-CPU.
Input matrix has last axis distributed amongst CPUs,
return is None on slaves, and the collected array on root.
The distribution can be uneven amongst CPUs. The list coords gives the
number of values for each CPU.
"""
a_xo = np.ascontiguousarray(a_xo)
if world.size == 1:
return a_xo
# All slaves send their piece to ``root``:
# There can be several sends before the corresponding receives
# are posted, so use syncronous send here
if world.rank != root:
world.ssend(a_xo, root, 112)
return None
# On root, put the subdomains from the slaves into the big array
# for the whole domain on root:
xshape = a_xo.shape[:-1]
Norb2 = sum(coords) # total number of orbital indices
a_xO = np.empty(xshape + (Norb2,), a_xo.dtype)
o = 0
for rank, norb in enumerate(coords):
if rank != root:
tmp_xo = np.empty(xshape + (norb,), a_xo.dtype)
world.receive(tmp_xo, rank, 112)
a_xO[..., o:o + norb] = tmp_xo
else:
a_xO[..., o:o + norb] = a_xo
o += norb
return a_xO
def collect_orbitals(a_xo, coords, root=0):
"""Collect array distributed over orbitals to root-CPU.
Input matrix has last axis distributed amongst CPUs,
return is None on slaves, and the collected array on root.
The distribution can be uneven amongst CPUs. The list coords gives the
number of values for each CPU.
"""
a_xo = np.ascontiguousarray(a_xo)
if world.size == 1:
return a_xo
# All slaves send their piece to ``root``:
# There can be several sends before the corresponding receives
# are posted, so use syncronous send here
if world.rank != root:
world.ssend(a_xo, root, 112)
return None
# On root, put the subdomains from the slaves into the big array
# for the whole domain on root:
xshape = a_xo.shape[:-1]
Norb2 = sum(coords) # total number of orbital indices
a_xO = np.empty(xshape + (Norb2,), a_xo.dtype)
o = 0
for rank, norb in enumerate(coords):
if rank != root:
tmp_xo = np.empty(xshape + (norb,), a_xo.dtype)
world.receive(tmp_xo, rank, 112)
a_xO[..., o:o + norb] = tmp_xo
else:
a_xO[..., o:o + norb] = a_xo
o += norb
return a_xO
def gatherv(m, N=None):
from gpaw.mpi import world, size, rank
if world.size == 1:
return m
ndim = m.ndim
if ndim == 2:
n, N = m.shape
assert n < N
M = np.zeros((N, N), dtype=complex)
elif ndim == 1:
n = m.shape[0]
M = np.zeros(N, dtype=complex)
else:
print 'Not Implemented'
XX
n_index = np.zeros(size, dtype=int)
world.all_gather(np.array([n]), n_index)
root = 0
if rank != root:
world.ssend(m, root, 112+rank)
else:
for irank, n in enumerate(n_index):
if irank == root:
if ndim == 2:
M[:n_index[0] :] = m
else:
M[:n_index[0]] = m
else:
n_start = n_index[0:irank].sum()
n_end = n_index[0:irank+1].sum()
if ndim == 2:
tmp_nN = np.zeros((n, N), dtype=complex)
world.receive(tmp_nN, irank, 112+irank)
M[n_start:n_end, :] = tmp_nN
else:
tmp_n = np.zeros(n, dtype=complex)
world.receive(tmp_n, irank, 112+irank)
M[n_start:n_end] = tmp_n
world.broadcast(M, root)
return M
def gatherv(m, N=None):
if world.size == 1:
return m
ndim = m.ndim
if ndim == 2:
n, N = m.shape
assert n < N
M = np.zeros((N, N), dtype=complex)
elif ndim == 1:
n = m.shape[0]
M = np.zeros(N, dtype=complex)
else:
print('Not Implemented')
XX
n_index = np.zeros(size, dtype=int)
world.all_gather(np.array([n]), n_index)
root = 0
if rank != root:
world.ssend(m, root, 112 + rank)
else:
for irank, n in enumerate(n_index):
if irank == root:
if ndim == 2:
M[:n_index[0]:] = m
else:
M[:n_index[0]] = m
else:
n_start = n_index[0:irank].sum()
n_end = n_index[0:irank + 1].sum()
if ndim == 2:
tmp_nN = np.zeros((n, N), dtype=complex)
world.receive(tmp_nN, irank, 112 + irank)
M[n_start:n_end, :] = tmp_nN
else:
tmp_n = np.zeros(n, dtype=complex)
world.receive(tmp_n, irank, 112 + irank)
M[n_start:n_end] = tmp_n
world.broadcast(M, root)
return M
