Repository: seung-lab/cloud-volume
Pull request #396 (2020-09-06 20:17:11)

perf(sharding): compressed morton code and gridpoints
- [x] compressed_morton_code
- [ ] gridpoints / xyzrange

#### feat(compressed_morton_code): accept arrays as input
For shard generation and reading large cutouts, several thousand codes need to be calculated. By allowing the entire array as input, we properly utilize numpy:
```python
grid_size = Vec(520,440,260)
gpts = np.array([np.random.randint(520, size=65536), np.random.randint(440, size=65536), np.random.randint(260, size=65536)]).T

%timeit for gridpoint in gpts: compressed_morton_code(gridpoint, grid_size)
# 4.2 s ± 168 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit compressed_morton_code(gpts, grid_size)
# 8.16 ms ± 350 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)




william-silversmith: Do you want to break this into two PRs so we can get the first speedup merged? This is v. cool.
nkemnitz: Feel free to merge. 

Patch:
diff --git a/cloudvolume/datasource/precomputed/image/__init__.py b/cloudvolume/datasource/precomputed/image/__init__.py
index af50c77..bbf0ebf 100644
--- a/cloudvolume/datasource/precomputed/image/__init__.py
+++ b/cloudvolume/datasource/precomputed/image/__init__.py
@@ -427,11 +427,10 @@ class PrecomputedImageSource(ImageSourceInterface):
     reader = sharding.ShardReader(self.meta, self.cache, spec)
 
     # 3. Gridpoints all within this one shard
-    gpts = lambda: gridpoints(aligned_bbox, self.meta.bounds(mip), chunk_size)
+    gpts = list(gridpoints(aligned_bbox, self.meta.bounds(mip), chunk_size))
+    morton_codes = compressed_morton_code(gpts, grid_size)
     all_same_shard = bool(reduce(lambda a,b: operator.eq(a,b) and a,
-      map(reader.get_filename,
-        map(lambda gpt: compressed_morton_code(gpt, grid_size), gpts())
-      )
+      map(reader.get_filename, morton_codes)
     ))
 
     if not all_same_shard:
@@ -440,8 +439,8 @@ class PrecomputedImageSource(ImageSourceInterface):
       )
 
     labels = {}
-    pt_anchor = next(gpts()) * chunk_size
-    for pt_abs in gpts():
+    pt_anchor = gpts[0] * chunk_size
+    for pt_abs, morton_code in zip(gpts, morton_codes):
       cutout_bbx = Bbox(pt_abs * chunk_size, (pt_abs + 1) * chunk_size)
 
       # Neuroglancer expects border chunks not to extend beyond dataset bounds
@@ -449,7 +448,6 @@ class PrecomputedImageSource(ImageSourceInterface):
       cutout_bbx -= pt_anchor
 
       chunk = img[ cutout_bbx.to_slices() ]
-      morton_code = compressed_morton_code(pt_abs, grid_size)
       labels[morton_code] = chunks.encode(chunk, self.meta.encoding(mip))
 
     shard_filename = reader.get_filename(first(labels.keys()))
diff --git a/cloudvolume/datasource/precomputed/image/common.py b/cloudvolume/datasource/precomputed/image/common.py
index 35404a5..da21fbd 100644
--- a/cloudvolume/datasource/precomputed/image/common.py
+++ b/cloudvolume/datasource/precomputed/image/common.py
@@ -68,18 +68,26 @@ def gridpoints(bbox, volume_bbox, chunk_size):
     yield Vec(x,y,z)
 
 def compressed_morton_code(gridpt, grid_size):
-  code = np.uint64(0)
-  num_bits = int(max([ np.ceil(np.log2(size)) for size in grid_size ]))
+  gridpt = np.asarray(gridpt, dtype=np.uint32)
+  single_input = False
+  if gridpt.ndim == 1:
+    gridpt = np.atleast_2d(gridpt)
+    single_input = True
+
+  code = np.zeros((gridpt.shape[0],), dtype=np.uint64)
+  num_bits = max(( math.ceil(math.log2(size)) for size in grid_size ))
   j = np.uint64(0)
   one = np.uint64(1)
 
   for i in range(num_bits):
     for dim in range(3):
       if 2 ** i <= grid_size[dim]:
-        bit = (((np.uint64(gridpt[dim]) >> np.uint64(i)) & one) << j)
+        bit = (((np.uint64(gridpt[:, dim]) >> np.uint64(i)) & one) << j)
         code |= bit
         j += one
-  
+
+  if single_input:
+    return code[0]
   return code
 
 def shade(dest_img, dest_bbox, src_img, src_bbox):
diff --git a/cloudvolume/datasource/precomputed/image/rx.py b/cloudvolume/datasource/precomputed/image/rx.py
index 3fd51d7..98f1385 100644
--- a/cloudvolume/datasource/precomputed/image/rx.py
+++ b/cloudvolume/datasource/precomputed/image/rx.py
@@ -55,16 +55,16 @@ def download_sharded(
 
   renderbuffer = np.zeros(shape=shape, dtype=meta.dtype, order=order)
 
-  gpts = gridpoints(full_bbox, bounds, chunk_size)
+  gpts = list(gridpoints(full_bbox, bounds, chunk_size))
 
   code_map = {}
-  for gridpoint in gpts:
-    zcurve_code = compressed_morton_code(gridpoint, grid_size)
+  morton_codes = compressed_morton_code(gpts, grid_size)
+  for gridpoint, morton_code in zip(gpts, morton_codes):
     cutout_bbox = Bbox(
       bounds.minpt + gridpoint * chunk_size,
       min2(bounds.minpt + (gridpoint + 1) * chunk_size, bounds.maxpt)
     )
-    code_map[zcurve_code] = cutout_bbox
+    code_map[morton_code] = cutout_bbox
 
   all_chunkdata = reader.get_data(list(code_map.keys()), meta.key(mip), progress=progress)
   for zcode, chunkdata in all_chunkdata.items():


Test patch:
diff --git a/test/test_sharding.py b/test/test_sharding.py
index 8b78bae..6be4636 100644
--- a/test/test_sharding.py
+++ b/test/test_sharding.py
@@ -63,6 +63,8 @@ def test_compressed_morton_code():
   assert cmc((0,0,7)) == 0b000100
   assert cmc((2,3,1)) == 0b011110
 
+  assert np.array_equal(cmc([(0,0,0), (1,0,1)]), [0b000000, 0b000101])
+
 def test_image_sharding_hash():
   spec = ShardingSpecification(
     type="neuroglancer_uint64_sharded_v1",