add correct roots for future work on more efficient host_row builder

ArgonneCPAC · Oct 17, 2023 · 50c01ee · 50c01ee
1 parent 3897ce5
commit 50c01ee
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/haccytrees/coretrees/__init__.py b/haccytrees/coretrees/__init__.py
@@ -1 +1,3 @@
-from .coretree_reader import corematrix_reader, coreforest2matrix
+from .coretree_reader import corematrix_reader
+
+__all__ = ("corematrix_reader",)
diff --git a/haccytrees/coretrees/coretree_reader.py b/haccytrees/coretrees/coretree_reader.py
@@ -172,11 +172,15 @@ def corematrix_reader(
             chunknum = 0
         if chunknum >= nchunks:
             print(f"invalid chunknum: {chunknum} needs to be smaller than {nchunks}")
-        start = roots[chunknum * chunksize]
+        startroot = chunknum * chunksize
+        endroot = min(nroots, (chunknum + 1) * chunksize)
+        start = roots[startroot]
         if chunknum == nchunks - 1:
             end = file_end
         else:
-            end = roots[(chunknum + 1) * chunksize]
+            end = roots[endroot]
+        roots = roots[startroot:endroot]
+        roots -= start
     else:
         start = roots[0]
         end = file_end