BUG: sparse: fix setdiag for matrices with missing diagonal entries (scipy#21792)

* BUG: sparse: fix setdiag for matrices with missing diagonal entries

We need to first set existing diagonal entries before adding
the missing ones, as inserting new values invalidates the computed
offsets.

* add tests for fix of csr.setdiag

---------

Co-authored-by: Dan Schult <[email protected]>

Author: sdrave

scipy/sparse/_compressed.py

@@ -970,17 +970,17 @@ def _setdiag(self, values, k):
             self.data[offsets] = x
             return
 
-        mask = (offsets <= -1)
+        mask = (offsets >= 0)
         # Boundary between csc and convert to coo
         # The value 0.001 is justified in gh-19962#issuecomment-1920499678
-        if mask.sum() < self.nnz * 0.001:
+        if self.nnz - mask.sum() < self.nnz * 0.001:
+            # replace existing entries
+            self.data[offsets[mask]] = x[mask]
             # create new entries
+            mask = ~mask
             i = i[mask]
             j = j[mask]
             self._insert_many(i, j, x[mask])
-            # replace existing entries
-            mask = ~mask
-            self.data[offsets[mask]] = x[mask]
         else:
             # convert to coo for _set_diag
             coo = self.tocoo()

scipy/sparse/tests/test_base.py

@@ -4272,6 +4272,13 @@ def test_scalar_idx_dtype(self):
         for x in [a, b, c, d, e, f]:
             x + x
 
+    def test_setdiag_csr(self):
+        # see gh-21791 setting mixture of existing and not when new_values < 0.001*nnz
+        D = self.dia_container(([np.arange(1002)], [0]), shape=(1002, 1002))
+        A = self.spcreator(D)
+        A.setdiag(5 * np.ones(A.shape[0]))
+        assert A[-1, -1] == 5
+
     def test_binop_explicit_zeros(self):
         # Check that binary ops don't introduce spurious explicit zeros.
         # See gh-9619 for context.
@@ -4441,6 +4448,13 @@ def test_scalar_idx_dtype(self):
         for x in [a, b, c, d, e, f]:
             x + x
 
+    def test_setdiag_csc(self):
+        # see gh-21791 setting mixture of existing and not when new_values < 0.001*nnz
+        D = self.dia_container(([np.arange(1002)], [0]), shape=(1002, 1002))
+        A = self.spcreator(D)
+        A.setdiag(5 * np.ones(A.shape[0]))
+        assert A[-1, -1] == 5
+
 
 TestCSC.init_class()