def matrix_elements(self,
other=None,
out=None,
symmetric=False,
cc=False,
operator=None,
result=None,
serial=False):
if out is None:
out = Matrix(len(self),
len(other or self),
dtype=self.dtype,
dist=(self.matrix.dist.comm, self.matrix.dist.rows,
self.matrix.dist.columns))
if other is None or isinstance(other, ArrayWaveFunctions):
assert cc
if other is None:
assert symmetric
operate_and_multiply(self, self.dv, out, operator, result)
elif not serial:
assert not symmetric
operate_and_multiply_not_symmetric(self, self.dv, out, other)
else:
self.multiply(self.dv, 'N', other, 'C', 0.0, out, symmetric)
else:
assert not cc
P_ani = {a: P_ni for a, P_ni in out.items()}
other.integrate(self.array, P_ani, self.kpt)
return out
def __init__(self,
nbands,
nproj_a,
atom_partition,
bcomm,
collinear=True,
spin=0,
dtype=float):
self.nproj_a = np.asarray(nproj_a)
self.atom_partition = atom_partition
self.bcomm = bcomm
self.collinear = collinear
self.spin = spin
self.nbands = nbands
self.indices = []
self.map = {}
I1 = 0
for a in self.atom_partition.my_indices:
ni = nproj_a[a]
I2 = I1 + ni
self.indices.append((a, I1, I2))
self.map[a] = (I1, I2)
I1 = I2
if not collinear:
I1 *= 2
self.matrix = Matrix(nbands, I1, dtype, dist=(bcomm, bcomm.size, 1))
if collinear:
self.myshape = self.matrix.array.shape
else:
self.myshape = (len(self.matrix.array), 2, I1 // 2)
def work_matrix_nn(self):
"""Get Matrix object for H, S, ..."""
if self._work_matrix_nn is None:
self._work_matrix_nn = Matrix(self.bd.nbands,
self.bd.nbands,
self.dtype,
dist=(self.bd.comm,
self.bd.comm.size))
return self._work_matrix_nn
def __init__(self, M, N, dtype, data, dist, collinear):
self.collinear = collinear
if not collinear:
N *= 2
if data is None or isinstance(data, np.ndarray):
self.matrix = Matrix(M, N, dtype, data, dist)
self.in_memory = True
else:
self.matrix = MatrixInFile(M, N, dtype, data, dist)
self.in_memory = False
self.comm = None
self.dtype = self.matrix.dtype
def matrix_elements(self,
other=None,
out=None,
symmetric=False,
cc=False,
operator=None,
result=None,
serial=False):
if other is None or isinstance(other, ArrayWaveFunctions):
if out is None:
out = Matrix(len(self),
len(other or self),
dtype=self.dtype,
dist=(self.matrix.dist.comm,
self.matrix.dist.rows,
self.matrix.dist.columns))
assert cc
if other is None:
assert symmetric
operate_and_multiply(self, self.dv, out, operator, result)
elif not serial:
assert not symmetric
operate_and_multiply_not_symmetric(self, self.dv, out, other)
elif self.dtype == complex:
self.matrix.multiply(self.dv, 'N', other.matrix, 'C', 0.0, out,
symmetric)
else:
self.matrix.multiply(2 * self.dv, 'N', other.matrix, 'T', 0.0,
out, symmetric)
if self.gd.comm.rank == 0:
correction = np.outer(self.matrix.array[:, 0],
other.matrix.array[:, 0])
if symmetric:
out.array -= 0.5 * self.dv * (correction +
correction.T)
else:
out.array -= self.dv * correction
else:
assert not cc
P_ani = {a: P_ni for a, P_ni in out.items()}
other.integrate(self.array, P_ani, self.kpt)
return out
def read_from_file(self):
"""Read wave functions from file into memory."""
matrix = Matrix(*self.matrix.shape,
dtype=self.dtype,
dist=self.matrix.dist)
# Read band by band to save memory
rows = matrix.dist.rows
blocksize = (matrix.shape[0] + rows - 1) // rows
for myn, psit_G in enumerate(matrix.array):
n = matrix.dist.comm.rank * blocksize + myn
if self.comm.rank == 0:
big_psit_G = self.array[n]
if big_psit_G.dtype == complex and self.dtype == float:
big_psit_G = big_psit_G.view(float)
elif big_psit_G.dtype == float and self.dtype == complex:
big_psit_G = np.asarray(big_psit_G, complex)
else:
big_psit_G = None
self._distribute(big_psit_G, psit_G)
self.matrix = matrix
self.in_memory = True
import time
import numpy as np
from gpaw.wavefunctions.arrays import UniformGridWaveFunctions
from gpaw.grid_descriptor import GridDescriptor
from gpaw.mpi import world, serial_comm
from gpaw.matrix import Matrix
S = world.size
B = 9
gd = GridDescriptor([32, 28, 112], [5, 5, 20], comm=serial_comm)
w = UniformGridWaveFunctions(B, gd, complex, dist=(world, S, 1))
#w.matrix.array.real[:] = np.arange(world.rank * B // S + 1,
# world.rank * B // S + 1 + B // S)[:, None]
w.array.imag[:] = np.random.uniform(-1, 1, w.array.shape)
w.array.real[:] = np.random.uniform(-1, 1, w.array.shape)
S0 = Matrix(B, B, complex, dist=(world, S, 1))
S0.array[:] = 42
S = Matrix(B, B, complex, dist=(world, S, 1))
S.array[:] = 42
t0 = time.time()
for i in range(1):
w.matrix_elements(w, symmetric=True, cc=True, out=S0)
t1 = time.time() - t0
#print(S.array, world.rank, S.array.shape)
t0 = time.time()
for i in range(1):
w.matrix_elements(symmetric=True, cc=True, out=S)
t2 = time.time() - t0
print(t1, t2)
#print(time.time() - t0)
#print(S.array, world.rank, S.array.shape)
import numpy as np
from gpaw.matrix import Matrix
from gpaw.mpi import world
N = 6
x = 0.01
A0 = Matrix(N, N, dist=(world, 1, 1), dtype=complex)
if world.rank == 0:
A0.array[:] = np.diag(np.arange(N) + 1)
A0.array += np.random.uniform(-x, x, (N, N))
A0.array += A0.array.conj().T
B = Matrix(N, N, data=A0.array.copy())
print(B.eigh(cc=True))
print(B.array)
A = Matrix(N, N, dist=(world, 2, 2, 2), dtype=complex)
A0.redist(A)
print(A.array)
print(A.eigh(cc=True, scalapack=(world, 2, 2, 2)))
print(world.rank, A.array)
A.redist(A0)
if world.rank == 0:
print(abs(A0.array) - abs(B.array))
print(A0.array / B.array)
import numpy as np
from gpaw.matrix import Matrix, matrix_matrix_multiply as mmm
from gpaw.mpi import world
N = 4
G = 7
# A0 = Matrix(N, N, dist=(world.new_communicator([0]), 1, 1))
A0 = Matrix(N, G, dist=(world, 1, 1), dtype=complex)
if world.rank == 0:
A0.array[:, 4:] = 1j
A0.array[:, :4] = np.diag(np.arange(N) + 1)
A = Matrix(N, G, dist=(world, world.size, 1), dtype=complex)
B = Matrix(N, G, dist=(world, world.size, 1), dtype=complex)
C = Matrix(N, N, dist=(world, world.size, 1), dtype=complex)
C0 = Matrix(N, N, dist=(world, 1, 1), dtype=complex)
A0.redist(A)
print(A.array)
A0.redist(B)
mmm(2.0, A, 'N', A, 'C', 0.0, C)
C.redist(C0)
print(C0.array)
C.array[:] = 777
mmm(2.0, A, 'N', A, 'C', 0.0, C, symmetric=True)
C.redist(C0)
print(C0.array)
N = 5
G = 7
A = Matrix(N, N, dist=(world, world.size, 1), dtype=complex)
B = Matrix(N, G, dist=(world, world.size, 1), dtype=complex)
C = Matrix(N, G, dist=(world, world.size, 1), dtype=complex)
from gpaw.matrix import Matrix
from gpaw.mpi import world
N = 6
if world.rank < 2:
comm = world.new_communicator([0, 1])
else:
comm = world.new_communicator([2, 3])
A0 = Matrix(N, N, dist=(comm, 2, 1))
A0.array[:] = world.rank
A = Matrix(N, N, dist=(world, 2, 2, 2))
A0.redist(A)
world.barrier()
print(A.array)
A0.array[:] = 117
A.redist(A0, 0)
print(A0.array)