Index: band_descriptor.py===================================================================--- band_descriptor.py (revision 5606)+++ band_descriptor.py (working copy)@@ -452,6 +452,197 @@ for q2 in range(Q): A_bnbn[(q1+q2)%B, :, q1] = A_qnn[q2]+#XXX vvvvvvvvvvvvvv++ def column_assembly(self, A_qnn, A_Nn, hermitian):+ """Assign all distributed sub-blocks pertaining from various rank to+ the relevant parts of a Hermitian or non-Hermitian column vector A_Nn.++ Parameters:++ A_qnn: ndarray+ Sub-blocks belonging to the specified rank.+ A_Nn: ndarray+ Full column vector in which to write contributions from sub-blocks.+ hermitian: bool+ Indicates whether A_Nn is part of a Hermitian matrix, in which+ case only the lower triangular part is assigned to.++ Note that the sub-block buffers are used for communicating across the+ band communicator, hence A_qnn will be altered during the assembly.+ """+ if self.comm.size == 1:+ self.columnwise_assign(A_qnn, A_Nn, 0, hermitian)+ return++ if self.rank == 0:+ for band_rank in range(self.comm.size):+ if band_rank > 0:+ self.comm.receive(A_qnn, band_rank, 13)+ self.columnwise_assign(A_qnn, A_Nn, band_rank, hermitian)+ else:+ self.comm.send(A_qnn, 0, 13)++ def columnwise_assign(self, A_qnn, A_Nn, band_rank, hermitian):+ """Assign the sub-blocks pertaining from a given rank to part of a+ blocked Hermitian or non-Hermitian matrix A_NN. This subroutine is+ used for column assembly.++ Parameters:++ A_qnn: ndarray+ Sub-blocks belonging to the specified rank.+ A_Nn: ndarray+ Full column vector in which to write contributions from sub-blocks.+ band_rank: int+ Communicator rank to which the sub-blocks belongs.+ hermitian: bool+ Indicates whether A_NN is a Hermitian matrix, in which+ case only the lower triangular part is assigned to.++ """+ if hermitian:+ self.triangular_columnwise_assign(A_qnn, A_Nn, band_rank)+ else:+ self.full_columnwise_assign(A_qnn, A_Nn, band_rank)++ def triangular_columnwise_assign(self, A_qnn, A_Nn, band_rank):+ """Assign the sub-blocks pertaining from a given rank to the lower+ triangular part of a Hermitian matrix A_NN. This subroutine is used+ for column assembly.++ Parameters:++ A_qnn: ndarray+ Sub-blocks belonging to the specified rank.+ A_Nn: ndarray+ Full column vector in which to write contributions from sub-blocks.+ band_rank: int+ Communicator rank to which the sub-blocks belongs.++ Note that a Hermitian matrix requires Q=B//2+1 blocks of M x M+ elements where B is the communicator size and M=N//B for N bands.+ """+ N = self.mynbands+ B = self.comm.size+ assert band_rank in xrange(B)++ if B == 1:+ # Only fill in the lower part+ mask = np.tri(N).astype(bool)+ A_NN[mask] = A_qnn.reshape((N,N))[mask]+ return++ # A_qnn[q2,myn1,myn2] on rank q1 is the q2'th overlap calculated+ # between <psi_n1| and A|psit_n2> where n1 <-> (q1,myn1) and + # n2 <-> ((q1+q2)%B,myn2) since we've sent/recieved q2 times.+ q1 = band_rank+ Q = B // 2 + 1+ if debug:+ assert A_qnn.shape == (Q,N,N)++ # Note that for integer inequalities, these relations are useful (X>0):+ # A*X > B <=> A > B//X ^ A*X <= B <=> A <= B//X++ if self.strided:+ A_nbn = A_NN.reshape((N, B, N))+ mask = np.empty((N,N), dtype=bool)+ for q2 in range(Q):+ # n1 = (q1+q2)%B + myn1*B ^ n2 = q1 + myn2*B+ #+ # We seek the lower triangular part i.e. n1 >= n2+ # <=> (myn2-myn1)*B <= (q1+q2)%B-q1+ # <=> myn2-myn1 <= dq//B+ dq = (q1+q2)%B-q1 # within ]-B; Q[ so dq//B is -1 or 0++ # Create mask for lower part of current block+ mask[:] = np.tri(N, N, dq//B)+ if debug:+ m1,m2 = np.indices((N,N))+ assert dq in xrange(-B+1,Q)+ assert (mask == (m1 >= m2 - dq//B)).all()++ # Copy lower part of A_qnn[q2] to its rightfull place+ A_nbn[:, (q1+q2)%B][mask] = A_qnn[q2][mask]++ # Negate the transposed mask to get complementary mask+ mask = ~mask.T++ # Copy upper part of Hermitian conjugate of A_qnn[q2]+ #A_nbnb[:, q1, :, (q1+q2)%B][mask] = A_qnn[q2].T.conj()[mask]+ raise NotImplementedError+ else:+ A_bnn = A_NN.reshape((B, N, N))++ # Optimization for the first block+ if q1 == 0:+ A_bnn[:Q] = A_qnn+ return++ for q2 in range(Q):+ # n1 = ((q1+q2)%B)*N + myn1 ^ n2 = q1*N + myn2+ #+ # We seek the lower triangular part i.e. n1 >= n2+ # <=> ((q1+q2)%B-q1)*N >= myn2-myn1+ # <=> myn2-myn1 <= dq*N+ # <=> entire block if dq > 0,+ # ... myn2 <= myn1 if dq == 0,+ # ... copy nothing if dq < 0+ if q1 + q2 < B:+ A_bnn[q1 + q2] = A_qnn[q2]+ else:+ #A_bnbn[q1, :, q1 + q2 - B] = A_qnn[q2].T.conj()+ raise NotImplementedError++ def full_columnwise_assign(self, A_qnn, A_Nn, band_rank):+ """Assign the sub-blocks pertaining from a given rank to the columns of+ non-Hermitian matrix A_Nn. This subroutine is used for column assembly.++ Parameters:++ A_qnn: ndarray+ Sub-blocks belonging to the specified rank.+ A_Nn: ndarray+ Full column vector, in which to write contributions from sub-blocks.+ band_rank: int+ Communicator rank to which the sub-blocks belongs.++ Note that a non-Hermitian matrix requires Q=B blocks of M x M+ elements where B is the communicator size and M=N//B for N bands.+ """+ N = self.mynbands+ B = self.comm.size+ assert band_rank in xrange(B)++ if B == 1:+ A_Nn[:] = A_qnn.reshape((N,N))+ return++ # A_qnn[q2,myn1,myn2] on rank q1 is the q2'th overlap calculated+ # between <psi_n1| and A|psit_n2> where n1 <-> (q1,myn1) and + # n2 <-> ((q1+q2)%B,myn2) since we've sent/recieved q2 times.+ q1 = band_rank+ Q = B+ if debug:+ assert A_qnn.shape == (Q,N,N)++ if self.strided:+ A_nbn = A_NN.reshape((N, B, N))+ for q2 in range(Q):+ A_nbn[:, (q1+q2)%B] = A_qnn[q2]+ else:+ A_bnn = A_Nn.reshape((B, N, N))++ # Optimization for the first block+ if q1 == 0:+ A_bnn[:Q] = A_qnn+ return++ for q2 in range(Q):+ A_bnn[(q1+q2)%B] = A_qnn[q2]++#XXX ^^^^^^^^^^^^^^+ def extract_block(self, A_NN, q1, q2): """Extract the sub-block pertaining from a given pair of ranks within the full matrix A_NN. Extraction may result in copies to assure unit