sd_writer.py

"""Write EXDF files based on EXtra-data SourceData objects.

Building on the EXDF DataFile API, SourceDataWriter builds
EXDf-structured HDF5 files based on a collection of EXtra-data
SourceData objects. It bypasses the DataCollection interface used in
extra_data.write to allow for more flexibility with data selection and
stricter validation when combining data from different locations.
"""


from functools import reduce
from itertools import accumulate
from logging import getLogger
from operator import or_

import numpy as np

from extra_data import FileAccess
from . import DataFile


log = getLogger('exdf.write.SourceDataWriter')


class SourceDataWriter:
    def get_data_format_version(self):
        """Determine the data format version."""
        return '1.3'

    def with_origin(self):
        """Determine whether to write INDEX/origin data."""
        return True

    def with_attrs(self):
        """Determine whether to write key attributes."""
        return True

    def chunk_instrument_data(self, source, key, orig_chunks):
        """Determine chunk size for INSTRUMENT key.

        Args:
            source (str): Source name.
            key (str): Key name.
            orig_chunks (tuple of int): Chunk size as found in input.

        Returns:
            (tuple of int): Chunk size to use for output.
        """

        return orig_chunks

    def copy_instrument_data(self, source, key, dest, data):
        """Copy INSTRUMENT data from input to output.

        Args:
            source (str): Source name.
            key (str): Key name.
            dset (h5py.Dataset): Destination dataset.
            data (ndarray): Source data.

        Returns:
            None
        """

        dest[:] = data

    def write_sequence(self, output_path, sources, sequence=0):
        """Write iterable of SourceData to file.

        Any train or key selection of each SourceData file is taken into
        account.

        Args:
            output_path (Path, str): Path to write to.
            sources (Iterable of extra_data.SourceData): Data sources.
            sequence (int, optional): Sequence number, 0 by default.

        Returns:
            None
        """

        with DataFile(output_path, 'w', driver='core') as f:
            self.write_base(f, sources, sequence)

            self.write_control(
                f, [sd for sd in sources if sd.is_control])

            self.write_instrument(
                f, [sd for sd in sources if sd.is_instrument])

    def write_base(self, f, sources, sequence):
        """Write METADATA, INDEX and source groups.

        Args:
            f (exdf.DataFile): Output file.
            sources (Iterable of extra_data.SourceData): Data sources.
            sequence (int, optional): Sequence number, 0 by default.

        Returns:
            None
        """

        train_ids, *index_dsets = get_index_root_data(sources)
        control_indices, instrument_indices = build_sources_index(sources)

        f.create_metadata(
            like=sources[0],
            sequence=sequence,
            data_format_version=self.get_data_format_version(),
            control_sources=control_indices.keys(),
            instrument_channels=[
                f'{source}/{channel}'
                for source, channels in instrument_indices.items()
                for channel in channels.keys()])
        f.create_dataset('METADATA/dataWriter', data=b'exdf-tools', shape=(1,))

        if not self.with_origin():
            index_dsets = (*index_dsets[:-1], None)

        f.create_index(train_ids, *index_dsets)

        for source, counts in control_indices.items():
            control_src = f.create_control_source(source)
            control_src.create_index(len(train_ids), per_train=True)

        for source, channel_counts in instrument_indices.items():
            instrument_src = f.create_instrument_source(source)
            instrument_src.create_index(**channel_counts)

    def write_control(self, f, sources):
        """Write CONTROL and RUN data.

        This method assumes the source datasets already exist.

        Args:
            f (exdf.DataFile): Output file.
            sources (Iterable of extra_data.SourceData): Data sources,
                should not contain instrument sources.

        Returns:
            None
        """

        for sd in sources:
            source = f.source[sd.source]

            attrs = get_key_attributes(sd) if self.with_attrs() else {}
            run_data_leafs = {}

            for path, value in sd.run_values().items():
                key = path[:path.rfind('.')]
                leaf = path[path.rfind('.')+1:]

                run_data_leafs.setdefault(key, {})[leaf] = value

            # Write CONTROL keys and their RUN keys.
            for key in sd.keys(False):
                run_entry = run_data_leafs.pop(key, False)

                if run_entry:
                    run_entry = (run_entry['value'], run_entry['timestamp'])

                ctrl_values = sd[f'{key}.value'].ndarray()
                ctrl_timestamps = sd[f'{key}.timestamp'].ndarray()

                source.create_key(
                    key, values=ctrl_values, timestamps=ctrl_timestamps,
                    run_entry=run_entry, attrs=attrs.pop(key, None))

            # Write remaining RUN-only keys.
            for key, leafs in run_data_leafs.items():
                source.create_run_key(key, **leafs, attrs=attrs.pop(key, None))

            # Fill in the missing attributes for nodes.
            for path, attrs in attrs.items():
                source.run_key[path].attrs.update(attrs)
                source.key[path].attrs.update(attrs)

    def write_instrument(self, f, sources):
        """Write INSTRUMENT data.

        This method assumes the source datasets already exist.

        Args:
            f (exdf.DataFile): Output file.
            sources (Iterable of extra_data.SourceData): Data sources,
                should not contain control sources.

        Returns:
            None
        """

        for sd in sources:
            source = f.source[sd.source]

            attrs = get_key_attributes(sd) if self.with_attrs() else {}

            for key in sd.keys():
                kd = sd[key]

                shape = (kd.data_counts(labelled=False).sum(), *kd.entry_shape)
                chunks = self.chunk_instrument_data(
                    sd.source, key,
                    kd.files[0].file[kd.hdf5_data_path].chunks)

                source.create_key(
                    key, shape=shape, maxshape=(None,) + shape[1:],
                    chunks=chunks, dtype=kd.dtype, attrs=attrs.pop(key, None))

            for path, attrs in attrs.items():
                source.key[path].attrs.update(attrs)

        # Copy INSTRUMENT data.
        for sd in sources:
            source = f.source[sd.source]

            for key in sd.keys():
                # TODO: Copy by chunk / file if too large
                self.copy_instrument_data(sd.source, key, source.key[key],
                                          sd[key].ndarray())


def get_index_root_data(sources):
    """Get train ID, timestamp, flag and origin data."""

    # Collect train IDs for this sequence.
    train_ids = np.zeros(0, dtype=np.uint64)
    for sd in sources:
        train_ids = np.union1d(train_ids, sd.train_ids)

    # Collect input files by index keys (source / index_group).
    files_by_index_keys = {}
    for sd in sources:
        for key in sd.keys():
            kd = sd[key]
            index_key = f'{sd.source}/{kd.index_group}'

            if index_key not in files_by_index_keys:
                files_by_index_keys[index_key] = set(kd.source_file_paths)

    inp_files = reduce(or_, files_by_index_keys.values(), set())

    # {trainId: ({item: timestamp}, {item: flag}, {item: origin})
    root_data = {}

    for inp_file in inp_files:
        fa = FileAccess(inp_file)

        sel_trains, sel_rows, _ = np.intersect1d(fa.train_ids, train_ids,
                                                 return_indices=True)

        if 'INDEX/timestamp' in fa.file:
            sel_timestamps = np.array(fa.file['INDEX/timestamp'][sel_rows])
        else:
            sel_timestamps = np.zeros_like(sel_trains, dtype=np.uint64)

        # Missing INDEX/flag dataset handled on the EXtra-data side.
        sel_flag = fa.validity_flag[sel_rows]

        if 'INDEX/origin' in fa.file:
            sel_origin = np.array(fa.file['INDEX/origin'][sel_rows])
        else:
            sel_origin = np.ones_like(sel_trains, dtype=np.int32)
            sel_origin[sel_flag] = 0

        item = (fa.data_category, fa.aggregator)

        # Collect timestamp, flag and origin as a function of train ID
        # and item while ensuring there is no mismatch for the same
        # train ID and item.
        for train_id, timestamp, flag, origin in zip(
            sel_trains, sel_timestamps, sel_flag, sel_origin
        ):
            new_entry = (timestamp, flag, origin)
            old_entry = root_data \
                .setdefault(train_id, dict()) \
                .setdefault(item, new_entry)

            if old_entry != new_entry:
                raise ValueError(f'INDEX root data mismatch for item {item}: '
                                 f'{new_entry} (from train {train_id}) vs '
                                 f'{old_entry}')

    if not np.array_equal(sorted(root_data.keys()), train_ids):
        # This is not supposed to happen.
        raise ValueError('input files missing index data for selected trains')

    timestamps = []
    flags = []
    origins = []

    # Decide which value to choose for each train.
    for train_id, entries in root_data.items():
        values, counts = np.unique(
            np.fromiter(entries.values(), dtype=object, count=len(entries)),
            return_counts=True)

        if len(values) > 1:
            values_str = ', '.join([f'{c}x {v}' for v, c
                                    in zip(values, counts)])
            log.warn(f'INDEX root data mismatch across items on train '
                     f'{train_id}: {values_str}')

        entry = values[counts.argmax()]
        timestamps.append(entry[0])
        flags.append(entry[1])
        origins.append(entry[2])

    return train_ids, np.array(timestamps), np.array(flags), np.array(origins)


def build_sources_index(sources):
    """Build data count and offsets."""

    # Build sources and indices.
    control_indices = {}
    instrument_indices = {}

    for sd in sources:
        if sd.is_control:
            control_indices[sd.source] = sd.data_counts(labelled=False)
        else:
            instrument_indices[sd.source] = {
                grp: sd.data_counts(labelled=False, index_group=grp)
                for grp in sd.index_groups}

    return control_indices, instrument_indices


def get_key_attributes(sd):
    source_attrs = dict()

    def build_path(a, b):
        return f'{a}.{b}'

    for key in sd.keys(inc_timestamps=False):
        paths = accumulate(key.split('.'), func=build_path)

        for path in paths:
            if path in source_attrs:
                # Skip this path, already parent of some other key.
                continue

            for source_file in sd[key].source_file_paths:
                fa = FileAccess(source_file)
                hdf_path = f'{sd.section}/{sd.source}/{path.replace(".", "/")}'
                path_attrs = dict(fa.file[hdf_path].attrs)

                existing_attrs = source_attrs.setdefault(path, path_attrs)

                if existing_attrs is not path_attrs:
                    same_keys = existing_attrs.keys() == path_attrs.keys()
                    same_values = all([
                        (
                            np.array_equal(v1, v2)
                            if isinstance(v1, np.ndarray)
                            else v1 == v2
                        )
                        for (k, v1), v2
                        in zip(existing_attrs.items(), path_attrs.values())
                        if k != 'tableSize'])

                    if not same_keys or not same_values:
                        log.debug(f'Attributes for {sd.source}.{path} in '
                                  f'{source_file} are {path_attrs}, but got '
                                  f'{source_attrs[path]} from previous file')
                        raise ValueError(f'attribute mismatch on '
                                         f'{sd.source}.{path}')

    return source_attrs