Пример #1
0
    def initialize(self, wfs):
        self.timer = wfs.timer
        self.world = wfs.world
        self.kpt_comm = wfs.kd.comm
        self.band_comm = wfs.band_comm
        self.dtype = wfs.dtype
        self.bd = wfs.bd
        self.ksl = wfs.diagksl
        self.nbands = wfs.bd.nbands
        self.mynbands = wfs.bd.mynbands
        self.operator = wfs.matrixoperator

        if self.mynbands != self.nbands or self.operator.nblocks != 1:
            self.keep_htpsit = False

        if self.keep_htpsit:
            self.Htpsit_nG = wfs.empty(self.nbands)
            if use_mic:
                self.Htpsit_nG_mic = stream.bind(self.Htpsit_nG)
                stream.sync()

        # Preconditioner for the electronic gradients:
        self.preconditioner = wfs.make_preconditioner(self.blocksize)

        for kpt in wfs.kpt_u:
            if kpt.eps_n is None:
                kpt.eps_n = np.empty(self.mynbands)

        # Allocate arrays for matrix operator
        self.operator.allocate_arrays()

        self.initialized = True
    def initialize(self, wfs):
        self.timer = wfs.timer
        self.world = wfs.world
        self.kpt_comm = wfs.kd.comm
        self.band_comm = wfs.band_comm
        self.dtype = wfs.dtype
        self.bd = wfs.bd
        self.ksl = wfs.diagksl
        self.nbands = wfs.bd.nbands
        self.mynbands = wfs.bd.mynbands
        self.operator = wfs.matrixoperator

        if self.mynbands != self.nbands or self.operator.nblocks != 1:
            self.keep_htpsit = False

        if self.keep_htpsit:
            self.Htpsit_nG = wfs.empty(self.nbands)
            if use_mic:
                self.Htpsit_nG_mic = stream.bind(self.Htpsit_nG)
                stream.sync()

        # Preconditioner for the electronic gradients:
        self.preconditioner = wfs.make_preconditioner(self.blocksize)

        for kpt in wfs.kpt_u:
            if kpt.eps_n is None:
                kpt.eps_n = np.empty(self.mynbands)
        
        # Allocate arrays for matrix operator
        self.operator.allocate_arrays()

        self.initialized = True
Пример #3
0
 def initialize_from_lcao_coefficients(self, basis_functions, mynbands):
     for kpt in self.kpt_u:
         kpt.psit_nG = self.gd.zeros(self.bd.mynbands, self.dtype)
         basis_functions.lcao_to_grid(kpt.C_nM,
                                      kpt.psit_nG[:mynbands], kpt.q)
         kpt.C_nM = None
         if use_mic:
             kpt.psit_nG_mic = stream.bind(kpt.psit_nG)
             stream.sync()
Пример #4
0
 def initialize_from_lcao_coefficients(self, basis_functions, mynbands):
     for kpt in self.kpt_u:
         kpt.psit_nG = self.gd.zeros(self.bd.mynbands, self.dtype)
         basis_functions.lcao_to_grid(kpt.C_nM, kpt.psit_nG[:mynbands],
                                      kpt.q)
         kpt.C_nM = None
         if use_mic:
             kpt.psit_nG_mic = stream.bind(kpt.psit_nG)
             stream.sync()
Пример #5
0
    def allocate_arrays(self):
        ngroups = self.bd.comm.size
        mynbands = self.bd.mynbands
        dtype = self.dtype
        if ngroups > 1:
            self.A_qnn = np.zeros((self.Q, mynbands, mynbands), dtype)
        self.A_nn = self.bmd.zeros(dtype=dtype)
        if use_mic:
            self.A_nn_mic = stream.bind(self.A_nn)
            stream.sync()

        if ngroups == 1 and self.nblocks == 1:
            self.work1_xG = self.gd.empty(self.bd.mynbands, self.dtype)
            if use_mic:
                self.work1_xG_mic = stream.bind(self.work1_xG)
                stream.sync()
        else:
            self.work1_xG = self.gd.empty(self.X, self.dtype)
            self.work2_xG = self.gd.empty(self.X, self.dtype)
    def allocate_arrays(self):
        ngroups = self.bd.comm.size
        mynbands = self.bd.mynbands
        dtype = self.dtype
        if ngroups > 1:
            self.A_qnn = np.zeros((self.Q, mynbands, mynbands), dtype)
        self.A_nn = self.bmd.zeros(dtype=dtype)
        if use_mic:
            self.A_nn_mic = stream.bind(self.A_nn)
            stream.sync()

        if ngroups == 1 and self.nblocks == 1:
            self.work1_xG = self.gd.empty(self.bd.mynbands, self.dtype) 
            if use_mic:
                self.work1_xG_mic = stream.bind(self.work1_xG)
                stream.sync()
        else:
            self.work1_xG = self.gd.empty(self.X, self.dtype)
            self.work2_xG = self.gd.empty(self.X, self.dtype)
    def empty(self, n=(), dtype=float, global_array=False, pad=False, usemic=False):
        """Return new uninitialized 3D array for this domain.

        The type can be set with the ``dtype`` keyword (default:
        ``float``).  Extra dimensions can be added with ``n=dim``.  A
        global array spanning all domains can be allocated with
        ``global_array=True``."""

        array = self._new_array(n, dtype, False, global_array, pad)
        if usemic:
            oa = stream.bind(array)
            stream.sync()
            return oa
        else:
            return array
    def integrate(self, a_xg, b_yg=None,
                  global_integral=True, hermitian=False,
                  _transposed_result=None):
        """Integrate function(s) over domain.

        a_xg: ndarray
            Function(s) to be integrated.
        b_yg: ndarray
            If present, integrate a_xg.conj() * b_yg.
        global_integral: bool
            If the array(s) are distributed over several domains, then the
            total sum will be returned.  To get the local contribution
            only, use global_integral=False.
        hermitian: bool
            Result is hermitian.
        _transposed_result: ndarray
            Long story.  Don't use this unless you are a method of the
            MatrixOperator class ..."""
        
        xshape = a_xg.shape[:-3]
        
        if b_yg is None:
            # Only one array:
            result = a_xg.reshape(xshape + (-1,)).sum(axis=-1) * self.dv
            if global_integral:
                if result.ndim == 0:
                    result = self.comm.sum(result)
                else:
                    self.comm.sum(result)
            return result

        if isinstance(a_xg, mic.OffloadArray):
            # offload arrays have to be contiguous in any case
            A_xg = a_xg
            B_yg = b_yg
        else:
            A_xg = np.ascontiguousarray(a_xg.reshape((-1,) + a_xg.shape[-3:]))
            B_yg = np.ascontiguousarray(b_yg.reshape((-1,) + b_yg.shape[-3:]))

        if _transposed_result is None:
            result_yx = np.zeros((len(B_yg), len(A_xg)), A_xg.dtype)
        else:
            result_yx = _transposed_result
            global_integral = False

        if isinstance(a_xg, mic.OffloadArray):
            result_yx_mic = stream.bind(result_yx)
            stream.sync()
            # result_yx_mic.fillfrom(result_yx)
            # result_yx_mic.array[:] = result_yx[:]
            # result_yx_mic.update_device()

        if a_xg is b_yg:
            if isinstance(a_xg, mic.OffloadArray):
                # dsyrk performs badly in MIC so use dgemm here
                # mic_rk(self.dv, A_xg, 0.0, result_yx_mic)
                mic_gemm(self.dv, A_xg, A_xg, 0.0, result_yx_mic, 'c')
            else:
                rk(self.dv, A_xg, 0.0, result_yx)
        elif hermitian:
            if isinstance(a_xg, mic.OffloadArray):
                mic_r2k(self.dv, A_xg, B_yg, 0.0, result_yx_mic)
            else:
                r2k(0.5 * self.dv, A_xg, B_yg, 0.0, result_yx)
        else:
            if isinstance(a_xg, mic.OffloadArray):
                mic_gemm(self.dv, A_xg, B_yg, 0.0, result_yx_mic, 'c')
            else:
                gemm(self.dv, A_xg, B_yg, 0.0, result_yx, 'c')
        
        if isinstance(a_xg, mic.OffloadArray):
            result_yx_mic.update_host()
            stream.sync()

        if global_integral:
            self.comm.sum(result_yx)

        yshape = b_yg.shape[:-3]
        result = result_yx.T.reshape(xshape + yshape)
        
        if result.ndim == 0:
            return result.item()
        else:
            return result
Пример #9
0
    def orthonormalize(self, wfs, kpt, psit_nG=None):
        """Orthonormalizes the vectors a_nG with respect to the overlap.

        First, a Cholesky factorization C is done for the overlap
        matrix S_nn = <a_nG | S | a_nG> = C*_nn C_nn Cholesky matrix C
        is inverted and orthonormal vectors a_nG' are obtained as::

          psit_nG' = inv(C_nn) psit_nG
                    __
           ~   _   \    -1   ~   _
          psi (r) = )  C    psi (r)
             n     /__  nm     m
                    m

        Parameters
        ----------

        psit_nG: ndarray, input/output
            On input the set of vectors to orthonormalize,
            on output the overlap-orthonormalized vectors.
        kpt: KPoint object:
            k-point object from kpoint.py.
        work_nG: ndarray
            Optional work array for overlap matrix times psit_nG.
        work_nn: ndarray
            Optional work array for overlap matrix.

        """
        self.timer.start('Orthonormalize')
        if psit_nG is None:
            psit_nG = kpt.psit_nG
            if use_mic:
                psit_nG_mic = kpt.psit_nG_mic
        else:
            if use_mic:
                psit_nG_mic = stream.bind(psit_nG, update_device=False)
                stream.sync()

        P_ani = kpt.P_ani
        self.timer.start('projections')
        wfs.pt.integrate(psit_nG, P_ani, kpt.q)
        self.timer.stop('projections')

        # Construct the overlap matrix:
        operator = wfs.matrixoperator

        def S(psit_G):
            return psit_G

        def dS(a, P_ni):
            return np.dot(P_ni, wfs.setups[a].dO_ii)

        if use_mic:
            self.timer.start('calc_s_matrix')
            psit_nG_mic.update_device()
            stream.sync()
            S_nn = operator.calculate_matrix_elements(psit_nG_mic, P_ani, S,
                                                      dS)
            self.timer.stop('calc_s_matrix')
        else:
            self.timer.start('calc_s_matrix')
            S_nn = operator.calculate_matrix_elements(psit_nG, P_ani, S, dS)
            self.timer.stop('calc_s_matrix')

        orthonormalization_string = repr(self.ksl)
        self.timer.start(orthonormalization_string)
        #
        if extra_parameters.get('sic', False):
            #
            # symmetric Loewdin Orthonormalization
            tri2full(S_nn, UL='L', map=np.conj)
            nrm_n = np.empty(S_nn.shape[0])
            diagonalize(S_nn, nrm_n)
            nrm_nn = np.diag(1.0 / np.sqrt(nrm_n))
            S_nn = np.dot(np.dot(S_nn.T.conj(), nrm_nn), S_nn)
        else:
            #
            self.ksl.inverse_cholesky(S_nn)
        # S_nn now contains the inverse of the Cholesky factorization.
        # Let's call it something different:
        C_nn = S_nn
        del S_nn
        self.timer.stop(orthonormalization_string)

        self.timer.start('rotate_psi')
        if use_mic:
            operator.matrix_multiply(C_nn,
                                     psit_nG_mic,
                                     P_ani,
                                     out_nG=kpt.psit_nG_mic)
            kpt.psit_nG_mic.update_host()
            stream.sync()
            # kpt.psit_nG[:] = self.psit_nG_mic.array[:]
        else:
            operator.matrix_multiply(C_nn, psit_nG, P_ani, out_nG=kpt.psit_nG)
        self.timer.stop('rotate_psi')
        self.timer.stop('Orthonormalize')
Пример #10
0
    def orthonormalize(self, wfs, kpt, psit_nG=None):
        """Orthonormalizes the vectors a_nG with respect to the overlap.

        First, a Cholesky factorization C is done for the overlap
        matrix S_nn = <a_nG | S | a_nG> = C*_nn C_nn Cholesky matrix C
        is inverted and orthonormal vectors a_nG' are obtained as::

          psit_nG' = inv(C_nn) psit_nG
                    __
           ~   _   \    -1   ~   _
          psi (r) = )  C    psi (r)
             n     /__  nm     m
                    m

        Parameters
        ----------

        psit_nG: ndarray, input/output
            On input the set of vectors to orthonormalize,
            on output the overlap-orthonormalized vectors.
        kpt: KPoint object:
            k-point object from kpoint.py.
        work_nG: ndarray
            Optional work array for overlap matrix times psit_nG.
        work_nn: ndarray
            Optional work array for overlap matrix.

        """
        self.timer.start('Orthonormalize')
        if psit_nG is None:
            psit_nG = kpt.psit_nG
            if use_mic:
                psit_nG_mic = kpt.psit_nG_mic
        else:
            if use_mic:
                psit_nG_mic = stream.bind(psit_nG, update_device=False)
                stream.sync()

        P_ani = kpt.P_ani
        self.timer.start('projections')
        wfs.pt.integrate(psit_nG, P_ani, kpt.q)
        self.timer.stop('projections')

        # Construct the overlap matrix:
        operator = wfs.matrixoperator

        def S(psit_G):
            return psit_G
        
        def dS(a, P_ni):
            return np.dot(P_ni, wfs.setups[a].dO_ii)

        if use_mic:
            self.timer.start('calc_s_matrix')
            psit_nG_mic.update_device()
            stream.sync()
            S_nn = operator.calculate_matrix_elements(psit_nG_mic, P_ani, S, dS)
            self.timer.stop('calc_s_matrix')
        else:
            self.timer.start('calc_s_matrix')
            S_nn = operator.calculate_matrix_elements(psit_nG, P_ani, S, dS)
            self.timer.stop('calc_s_matrix')


        orthonormalization_string = repr(self.ksl)
        self.timer.start(orthonormalization_string)
        #
        if extra_parameters.get('sic', False):
            #
            # symmetric Loewdin Orthonormalization
            tri2full(S_nn, UL='L', map=np.conj)
            nrm_n = np.empty(S_nn.shape[0])
            diagonalize(S_nn, nrm_n)
            nrm_nn = np.diag(1.0/np.sqrt(nrm_n))
            S_nn = np.dot(np.dot(S_nn.T.conj(), nrm_nn), S_nn)
        else:
            #
            self.ksl.inverse_cholesky(S_nn)
        # S_nn now contains the inverse of the Cholesky factorization.
        # Let's call it something different:
        C_nn = S_nn
        del S_nn
        self.timer.stop(orthonormalization_string)

        self.timer.start('rotate_psi')
        if use_mic:
            operator.matrix_multiply(C_nn, psit_nG_mic, P_ani, out_nG=kpt.psit_nG_mic)
            kpt.psit_nG_mic.update_host()
            stream.sync()
            # kpt.psit_nG[:] = self.psit_nG_mic.array[:]
        else:
            operator.matrix_multiply(C_nn, psit_nG, P_ani, out_nG=kpt.psit_nG)
        self.timer.stop('rotate_psi')
        self.timer.stop('Orthonormalize')