Skip to content

Commit

Permalink
add coretree chunking, add LC intersection code
Browse files Browse the repository at this point in the history
  • Loading branch information
michaelbuehlmann committed Nov 13, 2024
1 parent c4238fd commit 89aba8b
Show file tree
Hide file tree
Showing 5 changed files with 443 additions and 5 deletions.
3 changes: 2 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -46,4 +46,5 @@ _build
_notebooks
.venv

tmp/
tmp/
test/
21 changes: 17 additions & 4 deletions haccytrees/coretrees/coretree_reader.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Mapping, Union, List
from typing import Mapping, Union, List, Optional

import h5py
import numba
Expand Down Expand Up @@ -146,9 +146,9 @@ def corematrix_reader(
filename: str,
simulation: Union[Simulation, str],
*,
nchunks: int = None,
chunknum: int = None,
include_fields: List[str] = None,
nchunks: Optional[int] = None,
chunknum: Optional[int] = None,
include_fields: Optional[List[str]] = None,
calculate_host_rows: bool = True,
calculate_secondary_host_row: bool = False,
):
Expand Down Expand Up @@ -185,6 +185,15 @@ def corematrix_reader(
start = roots[0]
end = file_end

# how many rows did we skip?
if start > 0:
with h5py.File(filename) as forest_file:
core_tag = forest_file["data"]["core_tag"][:start]
skipped_rows = _count_coreforest_rows(core_tag)
del core_tag
else:
skipped_rows = 0

# read data
with h5py.File(filename) as forest_file:
if include_fields is None:
Expand All @@ -206,4 +215,8 @@ def corematrix_reader(
calculate_host_rows=calculate_host_rows,
calculate_secondary_host_row=calculate_secondary_host_row,
)

forest_matrices["absolute_row_idx"] = np.arange(
skipped_rows, skipped_rows + len(forest_matrices["core_tag"])
)
return forest_matrices
Empty file.
Loading

0 comments on commit 89aba8b

Please sign in to comment.