From 960ecfae06ee6fc6f7c1a395867284ea015bdce2 Mon Sep 17 00:00:00 2001 From: Pol Febrer Date: Wed, 8 Mar 2023 12:32:35 +0100 Subject: [PATCH] Discovering orbital pairs on the fly using ptr --- sisl/physics/_compute_dm.py | 66 ++++++++++++++++++++----------------- sisl/physics/compute_dm.py | 12 +++---- 2 files changed, 40 insertions(+), 38 deletions(-) diff --git a/sisl/physics/_compute_dm.py b/sisl/physics/_compute_dm.py index 68c3a386e6..61993068c7 100644 --- a/sisl/physics/_compute_dm.py +++ b/sisl/physics/_compute_dm.py @@ -7,7 +7,7 @@ @cython.boundscheck(False) @cython.wraparound(False) -def add_cnc_diag_spin(state: complex_or_float[:, :], row_orbs: cython.int[:], col_orbs_uc: cython.int[:], +def add_cnc_diag_spin(state: complex_or_float[:, :], DM_ptr: cython.int[:], DM_col_uc: cython.int[:], occs: cython.floating[:], DM_kpoint: complex_or_float[:], occtol: float = 1e-9): """Adds the cnc contributions of all orbital pairs to the DM given a array of states. @@ -18,10 +18,10 @@ def add_cnc_diag_spin(state: complex_or_float[:, :], row_orbs: cython.int[:], co state: The coefficients of all eigenstates for this contribution. Array of shape (n_eigenstates, n_basisorbitals) - row_orbs: - The orbital row indices of the sparsity pattern. - Shape (nnz, ), where nnz is the number of nonzero elements in the sparsity pattern. - col_orbs_uc: + DM_ptr: + The pointer to row array of the sparse DM. + Shape (no + 1, ), where no is the number of orbitals in the unit cell. + DM_col_uc: The orbital col indices of the sparsity pattern, but converted to the unit cell. Shape (nnz, ), where nnz is the number of nonzero elements in the sparsity pattern. occs: @@ -37,11 +37,13 @@ def add_cnc_diag_spin(state: complex_or_float[:, :], row_orbs: cython.int[:], co i: cython.int u: cython.int v: cython.int - ipair: cython.int + + # Number of orbitals in the unit cell + no: cython.int = DM_ptr.shape[0] - 1 + ival: cython.int # Loop lengths n_wfs: cython.int = state.shape[0] - n_opairs: cython.int = row_orbs.shape[0] # Variable to store the occupation of each state occ: float @@ -54,17 +56,17 @@ def add_cnc_diag_spin(state: complex_or_float[:, :], row_orbs: cython.int[:], co if occ < occtol: continue - # The occupation is above the tolerance threshold, loop through all overlaping orbital pairs - for ipair in range(n_opairs): - # Get the orbital indices of this pair - u = row_orbs[ipair] - v = col_orbs_uc[ipair] - # Add the contribution of this eigenstate to the DM_{u,v} element - DM_kpoint[ipair] = DM_kpoint[ipair] + state[i, u] * occ * state[i, v].conjugate() + # Loop over all non zero elements in the sparsity pattern + for u in range(no): + for ival in range(DM_ptr[u], DM_ptr[u+1]): + v = DM_col_uc[ival] + # Add the contribution of this eigenstate to the DM_{u,v} element + DM_kpoint[ival] = DM_kpoint[ival] + state[i, u] * occ * state[i, v].conjugate() + @cython.boundscheck(False) @cython.wraparound(False) -def add_cnc_nc(state: cython.complex[:, :, :], row_orbs: cython.int[:], col_orbs_uc: cython.int[:], +def add_cnc_nc(state: cython.complex[:, :, :], DM_ptr: cython.int[:], DM_col_uc: cython.int[:], occs: cython.floating[:], DM_kpoint: cython.complex[:, :, :], occtol: float = 1e-9): """Adds the cnc contributions of all orbital pairs to the DM given a array of states. @@ -76,10 +78,10 @@ def add_cnc_nc(state: cython.complex[:, :, :], row_orbs: cython.int[:], col_orbs The coefficients of all eigenstates for this contribution. Array of shape (n_eigenstates, n_basisorbitals, 2), where the last dimension is the spin "up"/"down" dimension. - row_orbs: - The orbital row indices of the sparsity pattern. - Shape (nnz, ), where nnz is the number of nonzero elements in the sparsity pattern. - col_orbs_uc: + DM_ptr: + The pointer to row array of the sparse DM. + Shape (no + 1, ), where no is the number of orbitals in the unit cell. + DM_col_uc: The orbital col indices of the sparsity pattern, but converted to the unit cell. Shape (nnz, ), where nnz is the number of nonzero elements in the sparsity pattern. occs: @@ -96,14 +98,17 @@ def add_cnc_nc(state: cython.complex[:, :, :], row_orbs: cython.int[:], col_orbs i: cython.int u: cython.int v: cython.int - ipair: cython.int + ival: cython.int + + # Number of orbitals in the unit cell + no: cython.int = DM_ptr.shape[0] - 1 + # The spin box indices. Di: cython.int Dj: cython.int # Loop lengths n_wfs: cython.int = state.shape[0] - n_opairs: cython.int = row_orbs.shape[0] # Variable to store the occupation of each state occ: float @@ -115,14 +120,13 @@ def add_cnc_nc(state: cython.complex[:, :, :], row_orbs: cython.int[:], col_orbs # If the occupation is lower than the tolerance, skip the state if occ < occtol: continue - - # The occupation is above the tolerance threshold, loop through all overlaping orbital pairs - for ipair in range(n_opairs): - # Get the orbital indices of this pair - u = row_orbs[ipair] - v = col_orbs_uc[ipair] + + # Loop over all non zero elements in the sparsity pattern + for u in range(no): + for ival in range(DM_ptr[u], DM_ptr[u+1]): + v = DM_col_uc[ival] - # Add to spin box - for Di in range(2): - for Dj in range(2): - DM_kpoint[ipair, Di, Dj] = DM_kpoint[ipair, Di, Dj] + state[i, u, Di] * occ * state[i, v, Dj].conjugate() \ No newline at end of file + # Add to spin box + for Di in range(2): + for Dj in range(2): + DM_kpoint[ival, Di, Dj] = DM_kpoint[ival, Di, Dj] + state[i, u, Di] * occ * state[i, v, Dj].conjugate() \ No newline at end of file diff --git a/sisl/physics/compute_dm.py b/sisl/physics/compute_dm.py index e6b287de4e..c7be1ebd73 100644 --- a/sisl/physics/compute_dm.py +++ b/sisl/physics/compute_dm.py @@ -46,9 +46,7 @@ def compute_dm(bz: BrillouinZone, eigenstates: Optional[Sequence[EigenstateElect geom = H.geometry # Sparsity pattern information - row_orbs, col_orbs = H.nonzero() - col_orbs_uc = H.osc2uc(col_orbs) - col_isc = col_orbs // H.no + col_isc, col_uc = np.divmod(H._csr.col, H.no) sc_offsets = H.sc_off.dot(H.cell) # Initialize the density matrix using the sparsity pattern of the Hamiltonian. @@ -121,8 +119,8 @@ def compute_dm(bz: BrillouinZone, eigenstates: Optional[Sequence[EigenstateElect if DM.spin.is_diagonal: # Calculate the matrix elements contributions for this k point. - DM_kpoint = np.zeros(row_orbs.shape[0], dtype=k_eigs.state.dtype) - add_cnc_diag_spin(state, row_orbs, col_orbs_uc, occs, DM_kpoint, occtol=occtol) + DM_kpoint = np.zeros(DM.nnz, dtype=k_eigs.state.dtype) + add_cnc_diag_spin(state, H._csr.ptr, col_uc, occs, DM_kpoint, occtol=occtol) # Apply phases DM_kpoint = DM_kpoint * phases @@ -139,8 +137,8 @@ def compute_dm(bz: BrillouinZone, eigenstates: Optional[Sequence[EigenstateElect # Calculate the matrix elements contributions for this k point. For each matrix element # we allocate a 2x2 spin box. - DM_kpoint = np.zeros((row_orbs.shape[0], 2, 2), dtype=np.complex128) - add_cnc_nc(state, row_orbs, col_orbs_uc, occs, DM_kpoint, occtol=occtol) + DM_kpoint = np.zeros((DM.nnz, 2, 2), dtype=np.complex128) + add_cnc_nc(state, H._csr.ptr, col_uc, occs, DM_kpoint, occtol=occtol) # Apply phases DM_kpoint *= phases.reshape(-1, 1, 1)