from collections import defaultdict
from pathlib import Path
import fnmatch
import logging
from packaging.version import Version
import numpy as np
from extra_data import by_id
from extra_data.read_machinery import select_train_ids
from exdf.write import SourceDataWriter
class ReduceWriter(SourceDataWriter):
log = logging.getLogger('exdf.data_reduction.ReduceWriter')
def __init__(self, data, methods, scope, sequence_len=-1, version=None):
self._data = data
self._methods = methods
self._scope = scope
self._sequence_len = sequence_len
metadata = self._data.run_metadata()
input_version = Version(metadata.get('dataFormatVersion', '1.0'))
if input_version < Version('1.0'):
raise ValueError('Currently input files are required to be '
'EXDF-v1.0+')
if version == 'same':
version = input_version
else:
self._version = Version(version)
try:
self.run_number = int(metadata['runNumber'])
except KeyError:
raise ValueError('runNumber dataset required in input METADATA')
self._ops = sum(methods.values(), [])
if not self._ops:
self.log.warning('Sum of reduction methods yielded no operations '
'to apply')
self._sources = sorted(data.all_sources)
self._touched_sources = set()
# Only populated if trains/keys are selected/removed for sources.
self._custom_keys = {} # source -> set(<keys>)
self._custom_trains = {} # source -> list(<trains>)
self._custom_xtdf_masks = {} # source -> dict(train_id -> mask)
self._custom_xtdf_counts = {} # source -> ndarray
self._custom_rows = {} # source -> dict(train_id -> mask)
self._rechunked_keys = {} # (source, key) -> chunks
self._partial_copies = {} # (source, key) -> list(<regions>)
# TODO: Raise error if rechunking is overwritten!
# TODO: make partial copies a list of slices!
# Collect reductions resulting from operations.
for source_glob, in self._filter_ops('remove-sources'):
for source in fnmatch.filter(self._sources, source_glob):
self._touched_sources.add(source)
self._sources.remove(source)
for source_glob, key_glob in self._filter_ops('remove-keys'):
for source in fnmatch.filter(self._sources, source_glob):
self._touched_sources.add(source)
keys = self._custom_keys.setdefault(
source, set(self._data[source].keys()))
for key in fnmatch.filter(keys, key_glob):
keys.remove(key)
for source_glob, train_sel in self._filter_ops('select-trains'):
for source in fnmatch.filter(self._sources, source_glob):
self._touched_sources.add(source)
train_ids = self._custom_trains.setdefault(
source, list(self._data.train_ids))
self._custom_trains[source] = select_train_ids(
train_ids, train_sel)
for source_glob, index_group, train_sel, row_sel in self._filter_ops(
'select-rows'
):
for source in fnmatch.filter(self._sources, source_glob):
if index_group not in self._data[source].index_groups:
raise ValueError(f'{index_group} not index group of '
f'{source}')
self._touched_sources.add(source)
self._custom_rows.setdefault((source, index_group), {}).update(
self._get_row_masks(source, index_group,
train_sel, row_sel))
for source_glob, train_sel, row_sel in self._filter_ops('select-xtdf'):
for source in fnmatch.filter(self._sources, source_glob):
if not source.endswith(':xtdf'):
# Simply ignore matches without trailing :xtdf.
continue
if not self._is_xtdf_source(source):
# Raise exception if essentials are missing.
raise ValueError(f'{source} is not a valid XTDF source')
self._touched_sources.add(source)
self._custom_xtdf_masks.setdefault(source, {}).update(
self._get_row_masks(source, 'image', train_sel, row_sel))
if (
{x[0] for x in self._custom_rows.keys()} &
self._custom_xtdf_masks.keys()
):
raise ValueError('source may not be affected by both select-rows '
'and select-xtdf operations')
for source_glob, key_glob, chunking in self._filter_ops(
'rechunk-keys'
):
for source in fnmatch.filter(self._sources, source_glob):
if not self._data[source].is_instrument:
raise ValueError(
f'rechunking keys only supported for instrument '
f'sources, but {source_glob} matches '
f'{self._data[source].section}/{source}')
self._touched_sources.add(source)
keys = self._custom_keys.get(
source, set(self._data[source].keys()))
for key in fnmatch.filter(keys, key_glob):
old_chunking = self._rechunked_keys.setdefault(
(source, key), chunking)
if old_chunking != chunking:
raise ValueError(
f'reduction sequence yields conflicting chunks '
f'for {source}.{key}: {old_chunking}, {chunking}')
self._rechunked_keys[(source, key)] = chunking
for source_glob, key_glob, region in self._filter_ops('partial-copy'):
for source in fnmatch.filter(self._sources, source_glob):
self._touched_sources.add(source)
keys = self._custom_keys.get(
source, set(self._data[source].keys()))
for key in fnmatch.filter(keys, key_glob):
self._partial_copies.setdefault((source, key), []).append(
region)
if self._scope == 'sources':
self._sources = sorted(
self._touched_sources.intersection(self._sources))
elif self._scope == 'aggregators':
touched_aggregators = {self._data[source].aggregator
for source in self._touched_sources}
self._sources = sorted(
{source for source in self._sources
if (self._data[source].aggregator in touched_aggregators)})
if not self._sources:
raise ValueError('reduction sequence yields empty source '
'selection')
def _filter_ops(self, op):
return [args[1:] for args in self._ops if args[0] == op]
def _is_xtdf_source(self, source):
return self._data[source].keys() > {'header.pulseCount', 'image.data'}
def _get_row_masks(self, source, index_group, train_sel, row_sel):
train_ids = select_train_ids(
self._custom_trains.get(source, list(self._data.train_ids)),
train_sel)
counts = self._data[source].select_trains(by_id[train_ids]) \
.data_counts(index_group=index_group)
masks = {}
if isinstance(row_sel, slice):
for train_id, count in counts.items():
if count > 0:
masks[train_id] = np.zeros(count, dtype=bool)
masks[train_id][row_sel] = True
elif np.issubdtype(type(row_sel[0]), np.integer):
max_row = max(row_sel)
for train_id, count in counts.items():
if count == 0:
continue
elif max_row >= count:
raise ValueError(
f'row index exceeds data counts of train {train_id}')
masks[train_id] = np.zeros(count, dtype=bool)
masks[train_id][row_sel] = True
elif np.issubdtype(type(row_sel[0]), bool):
mask_len = len(row_sel)
for train_id, count in counts.items():
if count == 0:
continue
elif mask_len != counts.get(train_id, 0):
raise ValueError(
f'mask length mismatch for train {train_id}')
masks[train_id] = row_sel
else:
raise ValueError('unknown row mask format')
return masks
def write_collection(self, output_path):
outp_data = self._data.select([(s, '*') for s in self._sources])
# Collect all items (combination of data category and
# aggregator) and the sources they contain.
sources_by_item = defaultdict(list)
for source in self._sources:
sd = outp_data[source]
sources_by_item[(sd.data_category, sd.aggregator)].append(source)
for (data_category, aggregator), sources in sources_by_item.items():
self.write_item(
output_path, sources, f'{data_category}-{aggregator}',
dict(data_category=data_category, aggregator=aggregator))
def write_collapsed(self, output_path):
self.write_item(output_path, self._sources, 'COLLAPSED')
def write_voview(self, output_path):
raise NotImplementedError('voview output layout')
def write_item(self, output_path, source_names, name, filename_fields={}):
"""Write sources to a single item."""
# Select output data down to what's in this item both in terms
# of sources and trains (via require_any).
item_data = self._data.select({
s: self._custom_keys[s] if s in self._custom_keys else set()
for s in source_names
}, require_any=True)
# Switch to representation of SourceData objects for
# per-source tracking of trains.
item_sources = [item_data[source]
for source in item_data.all_sources]
# Tetermine input sequence length if no explicit value was given
# for output.
if self._sequence_len < 1:
sequence_len = max({
len(sd._get_first_source_file().train_ids)
for sd in item_sources
})
else:
sequence_len = self._sequence_len
# Apply custom train selections, if any.
for i, sd in enumerate(item_sources):
train_sel = self._custom_trains.get(sd.source, None)
if train_sel is not None:
item_sources[i] = sd.select_trains(by_id[train_sel])
# Find the union of trains across all sources as total
# trains for this item.
item_train_ids = np.zeros(0, dtype=np.uint64)
for sd in item_sources:
item_train_ids = np.union1d(
item_train_ids, sd.drop_empty_trains().train_ids)
num_trains = len(item_train_ids)
num_sequences = int(np.ceil(num_trains / sequence_len))
self.log.info(
f'{name} containing {len(item_sources)} sources with {num_trains} '
f'trains over {num_sequences} sequences')
for seq_no in range(num_sequences):
seq_slice = np.s_[
(seq_no * sequence_len):((seq_no + 1) * sequence_len)]
# Slice out the train IDs and timestamps for this sequence.
seq_train_ids = item_train_ids[seq_slice]
# Select item data down to what's in this sequence.
seq_sources = [sd.select_trains(by_id[seq_train_ids])
for sd in item_sources]
# Build explicit output path for this sequence.
seq_path = Path(str(output_path).format(
run=self.run_number, sequence=seq_no, **filename_fields))
self.log.debug(f'{seq_path.stem} containing {len(seq_sources)} '
f'sources with {len(seq_train_ids)} trains')
self.write_sequence(seq_path, seq_sources, seq_no)
# SourceDataWriter hooks.
def write_base(self, f, sources, sequence):
super().write_base(f, sources, sequence)
# Add reduction-specific METADATA
red_group = f.require_group('METADATA/reduction')
for name, method in self._methods.items():
ops = np.array([
'\t'.join([str(x) for x in op[:]]).encode('ascii')
for op in method
])
red_group.create_dataset(name, shape=len(method), data=ops,)
def get_data_format_version(self):
return str(self._version)
def with_origin(self):
return self._version >= Version('1.2')
def with_attrs(self):
return self._version >= Version('1.3')
def chunk_instrument_data(self, source, key, orig_chunks):
try:
chunks = list(self._rechunked_keys[source, key])
assert len(chunks) == len(orig_chunks)
for i, dim_len in enumerate(chunks):
if dim_len is None:
chunks[i] = orig_chunks[i]
if -1 in chunks:
chunks[chunks.index(-1)] = \
np.prod(orig_chunks) // -np.prod(chunks)
return tuple(chunks)
except KeyError:
return orig_chunks
def mask_instrument_data(self, source, index_group, train_ids, counts):
if source in self._custom_xtdf_masks and index_group == 'image':
custom_masks = self._custom_xtdf_masks[source]
elif (source, index_group) in self._custom_rows:
custom_masks = self._custom_rows[source, index_group]
else:
return # None efficiently selects all rows.
masks = []
for train_id, count_all in zip(train_ids, counts):
if train_id in custom_masks:
mask = custom_masks[train_id]
else:
mask = np.ones(count_all, dtype=bool)
masks.append(mask)
if source in self._custom_xtdf_masks:
# Sources are guaranteed to never use both XTDF and general
# row slicing. In the XTDF case, the new data counts for the
# image index group must be determined to be filled into
# the respective header field.
self._custom_xtdf_counts[source] = {
train_id: mask.sum() for train_id, mask
in zip(train_ids, masks) if mask.any()}
return masks
def copy_instrument_data(self, source, key, dest, train_ids, data):
if source in self._custom_xtdf_counts and key == 'header.pulseCount':
custom_counts = self._custom_xtdf_counts[source]
for i, train_id in enumerate(train_ids):
data[i] = custom_counts.get(train_id, data[i])
try:
regions = self._partial_copies[source, key]
except KeyError:
dest[:] = data
else:
for region in regions:
sel = (np.s_[:], *region)
dest[sel] = data[sel]