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 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 __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 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)
class Projections: 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) @property def array(self): if self.collinear: return self.matrix.array else: return self.matrix.array.reshape(self.myshape) def new(self, bcomm='inherit', nbands=None, atom_partition=None): if bcomm == 'inherit': bcomm = self.bcomm elif bcomm is None: bcomm = serial_comm return Projections( nbands or self.nbands, self.nproj_a, self.atom_partition if atom_partition is None else atom_partition, bcomm, self.collinear, self.spin, self.matrix.dtype) def items(self): for a, I1, I2 in self.indices: yield a, self.array[..., I1:I2] def __getitem__(self, a): I1, I2 = self.map[a] return self.array[..., I1:I2] def __contains__(self, a): return a in self.map def todicttttt(self): return dict(self.items()) def redist(self, atom_partition): P = self.new(atom_partition=atom_partition) arraydict = self.toarraydict() arraydict.redistribute(atom_partition) P.fromarraydict(arraydict) return P def xxx_old_redist(self, atom_partition): P = self.new(atom_partition=atom_partition) rank_a = atom_partition.rank_a P_In = self.collect_atoms(self.matrix) if self.atom_partition.comm.rank == 0: mynbands = P_In.shape[1] for rank in range(self.atom_partition.comm.size): nI = self.nproj_a[rank_a == rank].sum() if nI == 0: continue P2_nI = np.empty((mynbands, nI), P_In.dtype) I1 = 0 myI1 = 0 for a, ni in enumerate(self.nproj_a): I2 = I1 + ni if rank == rank_a[a]: myI2 = myI1 + ni P2_nI[:, myI1:myI2] = P_In[I1:I2].T myI1 = myI2 I1 = I2 if rank == 0: P.matrix.array[:] = P2_nI else: self.atom_partition.comm.send(P2_nI, rank) else: if P.matrix.array.size > 0: self.atom_partition.comm.receive(P.matrix.array, 0) return P def collect(self): if self.bcomm.size == 1: P = self.matrix else: P = self.matrix.new(dist=(self.bcomm, 1, 1)) self.matrix.redist(P) if self.bcomm.rank > 0: return None if self.atom_partition.comm.size == 1: return P.array P_In = self.collect_atoms(P) if P_In is not None: return P_In.T def toarraydict(self): shape = self.myshape[:-1] shapes = [shape + (nproj, ) for nproj in self.nproj_a] d = self.atom_partition.arraydict(shapes, self.matrix.array.dtype) for a, I1, I2 in self.indices: d[a][:] = self.array[..., I1:I2] # Blocks will be contiguous return d def fromarraydict(self, d): assert d.partition == self.atom_partition for a, I1, I2 in self.indices: self.array[..., I1:I2] = d[a] def collect_atoms(self, P): if self.atom_partition.comm.rank == 0: nproj = sum(self.nproj_a) P_In = np.empty((nproj, P.array.shape[0]), dtype=P.array.dtype) I1 = 0 myI1 = 0 for nproj, rank in zip(self.nproj_a, self.atom_partition.rank_a): I2 = I1 + nproj if rank == 0: myI2 = myI1 + nproj P_In[I1:I2] = P.array[:, myI1:myI2].T myI1 = myI2 else: self.atom_partition.comm.receive(P_In[I1:I2], rank) I1 = I2 return P_In else: for a, I1, I2 in self.indices: self.atom_partition.comm.send(P.array[:, I1:I2].T.copy(), 0) return None
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)