diff --git a/cloudvolume/lib.py b/cloudvolume/lib.py index df9afab78..69c43997d 100644 --- a/cloudvolume/lib.py +++ b/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)) diff --git a/test/test_lib.py b/test/test_lib.py index 0813598b6..8f1672e75 100644 --- a/test/test_lib.py +++ 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():