def main(seed=42, dtype=float):
ksl = BlacsBandLayouts(gd, bd, block_comm, dtype, mcpus, ncpus, blocksize)
nbands = bd.nbands
mynbands = bd.mynbands
# Diagonalize
# We would *not* create H_Nn in the real-space code this way.
# This is just for testing purposes.
# Note after MPI_Reduce, only meaningful information on gd masters
H_Nn = np.zeros((nbands, mynbands), dtype=dtype)
U_nN = np.empty((mynbands, nbands), dtype=dtype)
if ksl.Nndescriptor: # hack
scalapack_set(ksl.Nndescriptor, H_Nn, 0.1, 75.0, 'L')
else:
assert gd.comm.rank != 0
print("H_Nn")
parallelprint(world, H_Nn)
eps_n = np.zeros(bd.mynbands)
blacs_diagonalize(ksl, H_Nn, U_nN, eps_n)
print("U_nN")
parallelprint(world, U_nN)
print("eps_n")
parallelprint(world, eps_n)
# Inverse Cholesky
S_Nn = np.zeros((nbands, mynbands), dtype=dtype)
C_nN = np.empty((mynbands, nbands), dtype=dtype)
if ksl.Nndescriptor: # hack
scalapack_set(ksl.Nndescriptor, S_Nn, 0.1, 75.0, 'L')
else:
assert gd.comm.rank != 0
print("S_Nn")
parallelprint(world, S_Nn)
blacs_inverse_cholesky(ksl, S_Nn, C_nN)
print("C_nN")
parallelprint(world, C_nN)
def main(seed=42, dtype=float):
ksl = BlacsBandLayouts(gd, bd, dtype, mcpus, ncpus, blocksize)
nbands = bd.nbands
mynbands = bd.mynbands
# Diagonalize
# We would *not* create H_Nn in the real-space code this way.
# This is just for testing purposes.
# Note after MPI_Reduce, only meaningful information on gd masters
H_Nn = np.zeros((nbands, mynbands), dtype=dtype)
U_nN = np.empty((mynbands, nbands), dtype=dtype)
if ksl.Nndescriptor: # hack
scalapack_set(ksl.Nndescriptor, H_Nn, 0.1, 75.0, 'L')
else:
assert gd.comm.rank != 0
print "H_Nn"
parallelprint(world, H_Nn)
eps_n = np.zeros(bd.mynbands)
blacs_diagonalize(ksl, H_Nn, U_nN, eps_n)
print "U_nN"
parallelprint(world, U_nN)
print "eps_n"
parallelprint(world, eps_n)
# Inverse Cholesky
S_Nn = np.zeros((nbands, mynbands), dtype=dtype)
C_nN = np.empty((mynbands, nbands), dtype=dtype)
if ksl.Nndescriptor: # hack
scalapack_set(ksl.Nndescriptor, S_Nn, 0.1, 75.0, 'L')
else:
assert gd.comm.rank != 0
print "S_Nn"
parallelprint(world, S_Nn)
blacs_inverse_cholesky(ksl, S_Nn, C_nN)
print "C_nN"
parallelprint(world, C_nN)
# simulate state-parallelization=2 and
# domain-decomposition.prod=32
B = 2
D = 32
mb = 32
grid = BlacsGrid(world, B, D)
nbands = 500
nG = 80**3
nGdesc = grid.new_descriptor(nbands, nG, nbands // B, nG // D)
nndesc = grid.new_descriptor(nbands, nbands, mb, mb)
psit_nG = gen.rand(*nGdesc.shape)
A_nn = gen.rand(*nndesc.shape)
assert nGdesc.check(psit_nG)
assert nndesc.check(A_nn)
parallelprint(world, (A_nn.shape, nndesc.shape, nndesc.lld))
pblas_simple_gemm(nGdesc,
nGdesc,
nndesc,
psit_nG,
psit_nG,
A_nn,
transa='N',
transb='T')
import numpy as np from gpaw.blacs import BlacsGrid, parallelprint from gpaw.mpi import world, rank, size from gpaw.utilities.scalapack import pblas_simple_gemm gen = np.random.RandomState(42) # simulate state-parallelization=2 and # domain-decomposition.prod=32 B = 2 D = 32 mb = 32 grid = BlacsGrid(world, B, D) nbands = 500 nG = 80 ** 3 nGdesc = grid.new_descriptor(nbands, nG, nbands / B, nG / D) nndesc = grid.new_descriptor(nbands, nbands, mb, mb) psit_nG = gen.rand(*nGdesc.shape) A_nn = gen.rand(*nndesc.shape) assert nGdesc.check(psit_nG) assert nndesc.check(A_nn) parallelprint(world, (A_nn.shape, nndesc.shape, nndesc.lld)) pblas_simple_gemm(nGdesc, nGdesc, nndesc, psit_nG, psit_nG, A_nn, transa="N", transb="T")