Fix sparse grid tests

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).any(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]

src/sisl/_core/sparse_geometry.py

@@ -235,7 +235,7 @@ def _translate_atoms_sc(self, sc_translations):
         # Get the row and column of every element in the matrix
         rows, cols = self.nonzero()
 
-        n_rows = len(self)
+        n_rows = self.shape[0]
         is_atom = n_rows == self.na
 
         # Find out the unit cell indices for the columns, and the index of the supercell

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

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
                 )

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

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