Repository: scipy/scipy
File: scipy/sparse/_compressed.py (Python)

Code before:
        M, N = self._swap((M, N))
        i, j = self._swap((i, j))
        n_samples = x.size
        offsets = np.empty(n_samples, dtype=self.indices.dtype)
        ret = csr_sample_offsets(M, N, self.indptr, self.indices, n_samples,
                                 i, j, offsets)
        if ret == 1:
            # rinse and repeat
            self.sum_duplicates()
            csr_sample_offsets(M, N, self.indptr, self.indices, n_samples,
                               i, j, offsets)
        if -1 not in offsets:
            # only affects existing non-zero cells
            self.data[offsets] = x
            return

        is_existing = (offsets >= 0)
        N_new = len(is_existing) - np.count_nonzero(is_existing)

        # Boundary between csc and convert to coo
        # The value 0.001 is justified in gh-19962#issuecomment-1920499678
        if N_new < self.nnz * 0.001:
            # replace existing entries
            self.data[offsets[is_existing]] = x[is_existing]
            # create new entries
            is_new = ~is_existing
            self._insert_many(i[is_new], j[is_new], x[is_new])
        else:
            # convert to coo for _set_diag
            do_sort = self.has_sorted_indices
            coo = self.tocoo()
            coo._setdiag(values, k)
            arrays = coo._coo_to_compressed(self._swap)
            self.indptr, self.indices, self.data, _ = arrays
            # Sort the indices (like in _insert_many)
            if do_sort:
                self.has_sorted_indices = False  # force a sort
                self.sort_indices()

    def _prepare_indices(self, i, j):
        M, N = self._swap(self._shape_as_2d)

        def check_bounds(indices, bound):
            idx = indices.max()
            if idx >= bound:
                raise IndexError(f'index ({idx}) out of range (>= {bound})')
            idx = indices.min()
            if idx < -bound:
                raise IndexError(f'index ({idx}) out of range (< -{bound})')

        i = np.atleast_1d(np.asarray(i, dtype=self.indices.dtype)).ravel()

Review comment (performance):
This is going to mean we store the boolean array twice (once for `is_existing` and another time for `is_new`), which is 2x the RAM cost. I agree that the old "mask" naming wasn't super clear, but it did allow us to save some allocations by reusing the same storage for the inverted array.

To be most efficient, we can do the negation in-place while assigning it to a new name for clarity: `is_new = np.logical_not(is_existing, out=is_existing)`

That will make `is_new` a simple alias for the modified `is_existing` array, but since we don't use `is_existing` anywhere else that should be fine. To be extra safe, we can `del is_existing` afterward.

Code after:
        M, N = self._swap((M, N))
        i, j = self._swap((i, j))
        n_samples = x.size
        offsets = np.empty(n_samples, dtype=self.indices.dtype)
        ret = csr_sample_offsets(M, N, self.indptr, self.indices, n_samples,
                                 i, j, offsets)
        if ret == 1:
            # rinse and repeat
            self.sum_duplicates()
            csr_sample_offsets(M, N, self.indptr, self.indices, n_samples,
                               i, j, offsets)
        if -1 not in offsets:
            # only affects existing non-zero cells
            self.data[offsets] = x
            return

        is_existing = (offsets >= 0)
        N_new = len(is_existing) - np.count_nonzero(is_existing)

        # Boundary between csc and convert to coo
        # The value 0.001 is justified in gh-19962#issuecomment-1920499678
        if N_new < self.nnz * 0.001:
            # replace existing entries
            self.data[offsets[is_existing]] = x[is_existing]
            # create new entries
            is_new = np.logical_not(is_existing, out=is_existing)
            del is_existing
            self._insert_many(i[is_new], j[is_new], x[is_new])
        else:
            # convert to coo for _set_diag
            do_sort = self.has_sorted_indices
            coo = self.tocoo()
            coo._setdiag(values, k)
            arrays = coo._coo_to_compressed(self._swap)
            self.indptr, self.indices, self.data, _ = arrays
            # Sort the indices (like in _insert_many)
            if do_sort:
                self.has_sorted_indices = False  # force a sort
                self.sort_indices()

    def _prepare_indices(self, i, j):
        M, N = self._swap(self._shape_as_2d)

        def check_bounds(indices, bound):
            idx = indices.max()
            if idx >= bound:
                raise IndexError(f'index ({idx}) out of range (>= {bound})')
            idx = indices.min()
            if idx < -bound:
                raise IndexError(f'index ({idx}) out of range (< -{bound})')