Merge pull request #458 from ZettaAI/zetta-avoid-copy

performance: Edge initialization
CAVEconnectome · Sep 5, 2023 · 6484f41 · 6484f41
2 parents e9e9492 + 82b0e9c
commit 6484f41
Showing 1 changed file with 9 additions and 6 deletions.
diff --git a/pychunkedgraph/graph/edges/__init__.py b/pychunkedgraph/graph/edges/__init__.py
@@ -30,21 +30,24 @@ def __init__(
         areas: Optional[np.ndarray] = None,
         fake_edges=False,
     ):
-        self.node_ids1 = np.array(node_ids1, dtype=basetypes.NODE_ID)
-        self.node_ids2 = np.array(node_ids2, dtype=basetypes.NODE_ID)
+        self.node_ids1 = np.array(node_ids1, dtype=basetypes.NODE_ID, copy=False)
+        self.node_ids2 = np.array(node_ids2, dtype=basetypes.NODE_ID, copy=False)
         assert self.node_ids1.size == self.node_ids2.size
-        self._affinities = np.ones(len(self.node_ids1)) * DEFAULT_AFFINITY
-        self._areas = np.ones(len(self.node_ids1)) * DEFAULT_AREA
+
         self._as_pairs = None
         self._fake_edges = fake_edges
 
         if affinities is not None and len(affinities) > 0:
-            self._affinities = np.array(affinities, dtype=basetypes.EDGE_AFFINITY)
+            self._affinities = np.array(affinities, dtype=basetypes.EDGE_AFFINITY, copy=False)
             assert self.node_ids1.size == self._affinities.size
+        else:
+            self._affinities = np.full(len(self.node_ids1), DEFAULT_AFFINITY)
 
         if areas is not None and len(areas) > 0:
-            self._areas = np.array(areas, dtype=basetypes.EDGE_AREA)
+            self._areas = np.array(areas, dtype=basetypes.EDGE_AREA, copy=False)
             assert self.node_ids1.size == self._areas.size
+        else:
+            self._areas = np.full(len(self.node_ids1), DEFAULT_AREA)
 
     @property
     def affinities(self) -> np.ndarray: