Update operators.py

loganbvh · Oct 2, 2023 · 0aa706e · 0aa706e
1 parent 57c4a6f
commit 0aa706e
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/tdgl/finite_volume/operators.py b/tdgl/finite_volume/operators.py
@@ -41,14 +41,13 @@ def _spmatrix_set_many(spmatrix, i, j, x):
     i, j = spmatrix._swap(i, j)
     offsets, (i, j, M, N) = _get_spmatrix_offsets_cupy(spmatrix, i, j)
 
-    if -1 not in offsets.get():
+    mask = offsets > -1
+    spmatrix.data[offsets[mask]] = x[mask]
+
+    if mask.all():
         # only affects existing non-zero cells
-        spmatrix.data[offsets] = x
         return
 
-    # replace where possible
-    mask = offsets > -1
-    spmatrix.data[offsets[mask]] = x[mask]
     # only insertions remain
     mask = ~mask
     i = i[mask]