diff --git a/src/sisl/_core/geometry.py b/src/sisl/_core/geometry.py index 708c6f059b..440746b6f6 100644 --- a/src/sisl/_core/geometry.py +++ b/src/sisl/_core/geometry.py @@ -3694,13 +3694,17 @@ def within_inf( # infinite supercell indices return self.asc2uc(idx), xyz, isc - def _orbital_values(self, grid_shape: tuple[int, int, int]): + def _orbital_values( + self, grid_shape: tuple[int, int, int], truncate_with_nsc: bool = False + ): r"""Calculates orbital values for a given grid. Parameters ---------- grid_shape: the grid shape (i.e. resolution) in which to calculate the orbital values. + truncate_with_nsc: + if True, only consider atoms within the geometry's auxiliary cell. Notes ----- @@ -3733,6 +3737,11 @@ def _orbital_values(self, grid_shape: tuple[int, int, int]): # supercell indices. Here we revert that ISC -= np.floor(self.fxyz[IA]).astype(int32) + # Don't consider atoms that are outside of the geometry's auxiliary cell. + if truncate_with_nsc: + mask = (abs(ISC) <= self.nsc // 2).all(axis=1) + IA, XYZ, ISC = IA[mask], XYZ[mask], ISC[mask] + def xyz2spherical(xyz, offset): """Calculate the spherical coordinates from indices""" rx = xyz[:, 0] - offset[0] diff --git a/src/sisl/_sparse_grid.py b/src/sisl/_sparse_grid.py index 478d8a11bd..dd4698ba7a 100644 --- a/src/sisl/_sparse_grid.py +++ b/src/sisl/_sparse_grid.py @@ -361,6 +361,19 @@ def reduce_orbital_products( The output grid. If ``out`` is not ``None``, it will be the same as ``out``. Otherwise it will be a newly created array. """ + if weights.shape[0] != self.geometry.no: + raise ValueError( + f"Mismatch in weights array shape: Number of unit cell orbitals is {weights.shape[0]}, while orbital values are stored for {self.geometry.no} orbitals." + ) + elif weights.shape[1] != self.geometry.no_s: + raise ValueError( + f"Mismatch in weights array shape: The number of unit cell orbitals ({self.geometry.no})" + f" is correct, but supercell orbitals in the weights array is {weights.shape[1]}," + f" while it should be {self.geometry.no_s}. It is likely that the weights array" + f" has been obtained from a matrix with the wrong number of auxiliary cells." + f" The correct number of auxiliary cells is {self.geometry.nsc}." + ) + csr = self._csr # Find out the reduced shape, and the reduce factor. The reduced factor is the number diff --git a/src/sisl/physics/densitymatrix.py b/src/sisl/physics/densitymatrix.py index c457d522a3..87b216e41f 100644 --- a/src/sisl/physics/densitymatrix.py +++ b/src/sisl/physics/densitymatrix.py @@ -715,7 +715,7 @@ def density( spinor=None, atol: float = 1e-7, eta: Optional[bool] = False, - method: Literal["pre-compute", "direct"] = "direct", + method: Literal["pre-compute", "direct"] = "pre-compute", **kwargs, ): r"""Expand the density matrix to the charge density on a grid @@ -758,9 +758,8 @@ def density( show a progressbar on stdout method: It determines if the orbital values are computed on the fly (direct) or they are all pre-computed - on the grid at the beginning(pre-compute). + on the grid at the beginning (pre-compute). Pre computing orbitals results in a faster computation, but it requires more memory. - Currently pre-computing has a bug that results in out-of-bounds, it will be fixed. Notes ----- @@ -772,6 +771,16 @@ def density( # same result. uc_dm = self.translate2uc() + if method == "pre-compute": + # Compute orbital values on the grid + psi_values = uc_dm.geometry._orbital_values(grid.shape) + + # Here we just set the nsc to whatever the psi values have. + # If the nsc is bigger in the DM, then some elements of the DM will be discarded. + # If the nsc is smaller in the DM, then the DM is just "padded" with 0s. + if not np.all(uc_dm.nsc == psi_values.geometry.nsc): + uc_dm.set_nsc(psi_values.geometry.nsc) + # Get the DM components with which we want to compute the density csr = uc_dm._csr if self.spin.kind > Spin.POLARIZED: @@ -806,7 +815,7 @@ def density( # Create the DM csr matrix. csrDM = csr_matrix( (DM, csr.col[idx], _ncol_to_indptr(csr.ncol)), - shape=(self.shape[:2]), + shape=(uc_dm.shape[:2]), dtype=DM.dtype, ) @@ -836,13 +845,7 @@ def density( csrDM = csr.tocsr(dim=0) if method == "pre-compute": - raise NotImplementedError( - "Currently a memory bug is happening, cannot be used until fixed, use the direct method" - ) try: - # Compute orbital values on the grid - psi_values = uc_dm.geometry._orbital_values(grid.shape) - psi_values.reduce_orbital_products( csrDM, uc_dm.lattice, out=grid.grid, **kwargs ) diff --git a/src/sisl/physics/tests/test_density_matrix.py b/src/sisl/physics/tests/test_density_matrix.py index d843cc28ea..313ab05d85 100644 --- a/src/sisl/physics/tests/test_density_matrix.py +++ b/src/sisl/physics/tests/test_density_matrix.py @@ -88,15 +88,7 @@ def func(D, ia, atoms, atoms_xyz): @pytest.fixture( scope="module", - params=[ - "direct", - pytest.param( - "pre-compute", - marks=pytest.mark.xfail( - reason="raises NotImplementedError due to seg-fault" - ), - ), - ], + params=["direct", "pre-compute"], ) def density_method(request): return request.param diff --git a/src/sisl/tests/test_sparse_grid.py b/src/sisl/tests/test_sparse_grid.py index 4328c45e55..c78774b7c4 100644 --- a/src/sisl/tests/test_sparse_grid.py +++ b/src/sisl/tests/test_sparse_grid.py @@ -17,6 +17,7 @@ def geometry(): orb = sisl.AtomicOrbital("2pzZ", (r, f)) geom = sisl.geom.graphene(orthogonal=True, atoms=sisl.Atom(6, orb)) + geom.set_nsc([3, 5, 1]) return geom @@ -92,7 +93,6 @@ def test_wavefunction_correct(H, grid_shape, psi_values, k): assert np.allclose(from_psi_grid.grid, wf_grid.grid) -@pytest.mark.skip(reason="bug in array indices, data-sc-off (sometimes)") @pytest.mark.parametrize( ["k", "ncoeffs"], [[(0, 0, 0), 1], [(0, 0, 0), 2], [(0.25, 0, 0), 1], [(0.25, 0, 0), 2]], @@ -149,7 +149,6 @@ def test_onthefly_reduction(geometry, psi_values, k, ncoeffs): assert np.allclose(reduced.reshape(post_reduced.shape), post_reduced) -@pytest.mark.skip(reason="bug in array indices, data-sc-off (sometimes)") def test_orbital_products(geometry): """Very simple tests to see if the orbital products are computed correctly""" geometry = geometry.copy() @@ -209,7 +208,6 @@ def test_orbital_products(geometry): assert np.allclose(dens.grid.ravel(), predicted) -@pytest.mark.skip(reason="bug in array indices, data-sc-off (sometimes)") def test_orbital_products_periodic(geometry): """Very simple tests to see if the orbital products are computed correctly when there are periodic conditions. @@ -250,7 +248,6 @@ def test_orbital_products_periodic(geometry): # We don't check for now. -@pytest.mark.skip(reason="bug in array indices, data-sc-off (sometimes)") @pytest.mark.parametrize("ncoeffs", [1, 2]) def test_orbital_products_onthefly_reduction(geometry, psi_values, ncoeffs): """Checks that the on the fly reduction produces the same