Skip to content
Snippets Groups Projects
Commit 32a660d9 authored by Philipp Schmidt's avatar Philipp Schmidt
Browse files

Merge branch 'feat/combine-entry-masks' into 'master'

Combine entry masks by logical AND rather than replace

See merge request !21
parents 76438f0f 9921f922
No related branches found
No related tags found
1 merge request!21Combine entry masks by logical AND rather than replace
......@@ -191,6 +191,8 @@ class ReduceWriter(SourceDataWriter):
return [args[1:] for args in self._ops if args[0] == op]
def _get_entry_masks(self, source, index_group, train_sel, entry_sel):
"""Generate bool vectors from any accepted selection."""
train_ids = select_train_ids(
self._custom_trains.get(source, list(self._data.train_ids)),
train_sel)
......@@ -466,8 +468,16 @@ class ReduceWriter(SourceDataWriter):
raise ReduceInitError(
f'{idx_group} not an index group of {source}')
self._custom_entry_masks.setdefault((source, idx_group), {}).update(
self._get_entry_masks(source, idx_group, train_sel, entry_sel))
masks = self._custom_entry_masks.setdefault((source, idx_group), {})
new_masks = self._get_entry_masks(
source, idx_group, train_sel, entry_sel)
for train_id, train_mask in new_masks.items():
if train_id in masks:
train_mask &= masks[train_id]
masks[train_id] = train_mask
self.log.debug(f'Applying entry selection to {source}, {idx_group}')
@apply_by_source('select-xtdf')
......@@ -477,8 +487,16 @@ class ReduceWriter(SourceDataWriter):
f'Ignoring non-XTDF source {source} based on name')
return False
self._custom_xtdf_masks.setdefault(source, {}).update(
self._get_entry_masks(source, 'image', train_sel, entry_sel))
masks = self._custom_xtdf_masks.setdefault(source, {})
new_masks = self._get_entry_masks(
source, 'image', train_sel, entry_sel)
for train_id, train_mask in new_masks.items():
if train_id in masks:
train_mask &= masks[train_id]
masks[train_id] = train_mask
self.log.debug(f'Applying XTDF selection to {source}')
@apply_by_key('rechunk-keys')
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment