Remove extra copy

flatironinstitute · Feb 28, 2024 · 2e6dfb7 · 2e6dfb7
1 parent 4888d24
commit 2e6dfb7
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 4 deletions.
diff --git a/inferelator_velocity/denoise.py b/inferelator_velocity/denoise.py
@@ -1,3 +1,5 @@
+import numpy as np
+
 from inferelator_velocity.utils.noise2self import (
     _dist_to_row_stochastic,
     dot
@@ -24,10 +26,13 @@ def denoise(
             f"run global_graph() first"
         )
 
-    data.layers[output_layer] = dot(
+    data.layers[output_layer] = np.zeros(lref.shape, dtype=np.float32)
+
+    dot(
         _dist_to_row_stochastic(data.obsp[graph_key]),
         lref,
-        dense=dense
+        dense=dense,
+        out=data.layers[output_layer]
     )
 
     return data
diff --git a/inferelator_velocity/utils/math.py b/inferelator_velocity/utils/math.py
@@ -19,14 +19,18 @@
         f"{str(err)}"
     )
 
-    def dot(x, y, dense=False, cast=False):
+    def dot(x, y, dense=False, cast=False, out=None):
 
         z = x @ y
 
         if dense and sps.issparse(z):
             z = z.A
 
-        return z
+        if out is not None:
+            out[:] = z
+            return out
+        else:
+            return z
 
 
 def scalar_projection(