wip+perf: faster xyzrange

cloudvolume/lib.py

@@ -208,10 +208,32 @@ def vectorize():
     pt = Vec(0,0,0)
     for z,y,x in product(*zyxranges):
       pt.x, pt.y, pt.z = min(x, end_vec[0]), min(y, end_vec[1]), min(z, end_vec[2])
-      yield pt
+      yield pt.clone()
 
   return vectorize()
 
+def xyzrange_np(start_vec, end_vec=None, stride_vec=(1,1,1)):
+  if end_vec is None:
+    end_vec = start_vec
+    start_vec = (0,0,0)
+
+  start_vec = np.array(start_vec, dtype=np.int64)
+  end_vec = np.array(end_vec, dtype=np.int64)  
+
+  delta_vec = end_vec - start_vec
+  if not np.any(delta_vec):
+    return np.array([[]], dtype=np.int64)
+
+  axes = []
+  for axis in range(start_vec.size):
+    arr = list(range(start_vec[axis], end_vec[axis], stride_vec[axis]))
+    if arr[-1] < end_vec[axis] - 1:
+      arr.append(end_vec[axis])
+    axes.append(arr)
+
+  return np.array(np.meshgrid(*axes)).T.reshape(-1,3)
+
+
 def map2(fn, a, b):
   assert len(a) == len(b), "Vector lengths do not match: {} (len {}), {} (len {})".format(a[:3], len(a), b[:3], len(b))
 

test/test_lib.py

@@ -3,7 +3,24 @@
 import numpy as np
 
 import cloudvolume.lib as lib
-from cloudvolume.lib import Bbox, Vec
+from cloudvolume.lib import Bbox, Vec, xyzrange, xyzrange_np
+
+def test_xyzrange():
+  def xyz(*args):
+    return np.array(list(xyzrange_np(*args)))
+
+  assert list(xyzrange((0,0,0))) == []
+  assert list(xyzrange((1,0,0))) == []
+  assert np.all(xyz((1,1,1)) == [Vec(0,0,0)])
+  assert np.all(xyz((2,1,1)) == [Vec(0,0,0), Vec(1,0,0)])
+  assert np.all(xyz((2,2,2)) == [
+    Vec(0,0,0), Vec(1,0,0),
+    Vec(0,1,0), Vec(1,1,0),
+    Vec(0,0,1), Vec(1,0,1),
+    Vec(0,1,1), Vec(1,1,1),
+  ])
+  assert np.all(xyz((2,1,1), (3,2,2)) == [Vec(2,1,1)])
+  assert np.all(xyz((2,1,1), (5,2,2), (2,1,1)) == [Vec(2,1,1), Vec(4,1,1)])
 
 def test_divisors():