<commit_before>"""Data accessor class for HDF5 files produced by KAT-7 correlator."""

import logging

import numpy as np
import h5py
import katpoint

from .dataset import DataSet, WrongVersion, BrokenFile, Subarray, \
                     SpectralWindow, DEFAULT_VIRTUAL_SENSORS, _robust_target
from .sensordata import SensorData, SensorCache
from .categorical import CategoricalData, sensor_to_categorical
from .lazy_indexer import LazyIndexer

logger = logging.getLogger('katfile.h5datav2')

# Simplify the scan activities to derive the basic state of the antenna (slewing, scanning, tracking, stopped)
SIMPLIFY_STATE = {'scan_ready': 'slew', 'scan': 'scan', 'scan_complete': 'scan', 'track': 'track', 'slew': 'slew'}

SENSOR_PROPS = {
    '*activity' : {'greedy_values' : ('slew', 'stop'), 'initial_value' : 'slew',
                   'transform' : lambda act: SIMPLIFY_STATE.get(act, 'stop')},
    '*target' : {'initial_value' : '', 'transform' : _robust_target},
    'RFE/center-frequency-hz' : {'categorical' : True},
    'RFE/rfe7.lo1.frequency' : {'categorical' : True},
    'Observation/label' : {'initial_value' : '', 'transform' : str, 'allow_repeats' : True}
}

def _calc_azel(cache, name, ant):
    """Calculate virtual (az, el) sensors from actual ones in sensor cache."""
    real_sensor = 'Antennas/%s/%s' % (ant, 'pos.actual-scan-azim' if name.endswith('az') else 'pos.actual-scan-elev')
    cache[name] = sensor_data = katpoint.deg2rad(cache.get(real_sensor))
    return sensor_data

VIRTUAL_SENSORS = dict(DEFAULT_VIRTUAL_SENSORS)
VIRTUAL_SENSORS.update({'Antennas/{ant}/az' : _calc_azel, 'Antennas/{ant}/el' : _calc_azel})

#--------------------------------------------------------------------------------------------------
#--- Utility functions
#--------------------------------------------------------------------------------------------------

def get_single_value(group, name):
    """Return single value from attribute or dataset with given name in group.

    If `name` is an attribute of the HDF5 group `group`, it is returned,
    otherwise it is interpreted as an HDF5 dataset of `group` and the last value
    of `name` is returned. This is meant to retrieve static configuration values
    that potentially get set more than once during capture initialisation, but
    then does not change during actual capturing.

    Parameters
    ----------
    group : :class:`h5py.Group` object
        HDF5 group to query
    name : string
        Name of HDF5 attribute or dataset to query

    Returns
    -------
    value : object
        Attribute or last value of dataset

    """
    return group.attrs[name] if name in group.attrs else group[name].value[-1]

#--------------------------------------------------------------------------------------------------
#--- CLASS :  H5DataV2
#--------------------------------------------------------------------------------------------------

class H5DataV2(DataSet):
    """Load HDF5 format version 2 file produced by KAT-7 correlator.

    For more information on attributes, see the :class:`DataSet` docstring.

    Parameters
    ----------
    filename : string
        Name of HDF5 file
    ref_ant : string, optional
        Name of reference antenna, used to partition data set into scans
        (default is first antenna in use)
    time_offset : float, optional
        Offset to add to all correlator timestamps, in seconds

    Attributes
    ----------
    file : :class:`h5py.File` object
        Underlying HDF5 file, exposed via :mod:`h5py` interface

    """
    def __init__(self, filename, ref_ant='', time_offset=0.0):
        DataSet.__init__(self, filename, ref_ant, time_offset)

        # Load file
        self.file = f = h5py.File(filename, 'r')

        # Only continue if file is correct version and has been properly augmented
        self.version = f.attrs.get('version', '1.x')
        if not self.version.startswith('2.'):
            raise WrongVersion("Attempting to load version '%s' file with version 2 loader" % (self.version,))
        if not 'augment_ts' in f.attrs:
            raise BrokenFile('HDF5 file not augmented - please run k7_augment.py (provided by katcapture package)')

        # Load main HDF5 groups
        data_group, sensors_group, config_group = f['Data'], f['MetaData/Sensors'], f['MetaData/Configuration']
        markup_group = f['Markup']
        # Get observation script attributes, with defaults
        script_attrs = config_group['Observation'].attrs
        self.observer = script_attrs.get('script_observer', '')
        self.description = script_attrs.get('script_description', '')
        self.experiment_id = script_attrs.get('script_experiment_id', '')

        # ------ Extract timestamps ------

        self.dump_period = get_single_value(config_group['Correlator'], 'int_time')
        # Obtain visibility data and timestamps
        self._vis = data_group['correlator_data']
        self._timestamps = data_group['timestamps']
        num_dumps = len(self._timestamps)
        # Discard the last sample if the timestamp is a duplicate (caused by stop packet in k7_capture)
        num_dumps = (num_dumps - 1) if num_dumps > 1 and (self._timestamps[-1] == self._timestamps[-2]) else num_dumps
        # Estimate timestamps by assuming they are uniformly spaced (much quicker than loading them from file).
        # This is useful for the purpose of segmenting data set, where accurate timestamps are not that crucial.
        # The real timestamps are still loaded when the user explicitly asks for them.
        # Do quick test for uniform spacing of timestamps (necessary but not sufficient).
        if abs((self._timestamps[num_dumps - 1] - self._timestamps[0]) / self.dump_period + 1 - num_dumps) < 0.01:
            # Estimate the timestamps as being uniformly spaced
            data_timestamps = self._timestamps[0] + self.dump_period * np.arange(num_dumps)
        else:
            # Load the real timestamps instead and warn the user, as this is anomalous
            data_timestamps = self._timestamps[:num_dumps]
            expected_dumps = (self._timestamps[num_dumps - 1] - self._timestamps[0]) / self.dump_period + 1
            logger.warning(("Irregular timestamps detected in file '%s':"
                           "expected %.3f dumps based on dump period and start/end times, got %d instead") %
                           (filename, expected_dumps, num_dumps))
        # Move timestamps from start of each dump to the middle of the dump
        data_timestamps += 0.5 * self.dump_period + self.time_offset
        if data_timestamps[0] < 1e9:
            logger.warning("File '%s' has invalid first correlator timestamp (%f)" % (filename, data_timestamps[0],))
        self._time_keep = np.ones(num_dumps, dtype=np.bool)
        self.start_time = katpoint.Timestamp(data_timestamps[0] - 0.5 * self.dump_period)
        self.end_time = katpoint.Timestamp(data_timestamps[-1] + 0.5 * self.dump_period)

        # ------ Extract sensors ------

        # Populate sensor cache with all HDF5 datasets below sensor group that fit the description of a sensor
        cache = {}
        def register_sensor(name, obj):
            """A sensor is defined as a non-empty dataset with expected dtype."""
            if isinstance(obj, h5py.Dataset) and obj.shape != () and obj.dtype.names == ('timestamp','value','status'):
                cache[name] = SensorData(obj, name)
        sensors_group.visititems(register_sensor)
        # Use estimated data timestamps for now, to speed up data segmentation<commit_msg>Rename pedestal sensors from now on to fall under their corresponding antenna groups in v2 files. This matches the situation in v1 files and will soon be the standard anyway, as pedestal proxies have been merged into antenna proxies in trunk of November 2011. It also simplifies user scripts.<commit_after>"""Data accessor class for HDF5 files produced by KAT-7 correlator."""

import logging

import numpy as np
import h5py
import katpoint

from .dataset import DataSet, WrongVersion, BrokenFile, Subarray, \
                     SpectralWindow, DEFAULT_VIRTUAL_SENSORS, _robust_target
from .sensordata import SensorData, SensorCache
from .categorical import CategoricalData, sensor_to_categorical
from .lazy_indexer import LazyIndexer

logger = logging.getLogger('katfile.h5datav2')

# Simplify the scan activities to derive the basic state of the antenna (slewing, scanning, tracking, stopped)
SIMPLIFY_STATE = {'scan_ready': 'slew', 'scan': 'scan', 'scan_complete': 'scan', 'track': 'track', 'slew': 'slew'}

SENSOR_PROPS = {
    '*activity' : {'greedy_values' : ('slew', 'stop'), 'initial_value' : 'slew',
                   'transform' : lambda act: SIMPLIFY_STATE.get(act, 'stop')},
    '*target' : {'initial_value' : '', 'transform' : _robust_target},
    'RFE/center-frequency-hz' : {'categorical' : True},
    'RFE/rfe7.lo1.frequency' : {'categorical' : True},
    'Observation/label' : {'initial_value' : '', 'transform' : str, 'allow_repeats' : True}
}

def _calc_azel(cache, name, ant):
    """Calculate virtual (az, el) sensors from actual ones in sensor cache."""
    real_sensor = 'Antennas/%s/%s' % (ant, 'pos.actual-scan-azim' if name.endswith('az') else 'pos.actual-scan-elev')
    cache[name] = sensor_data = katpoint.deg2rad(cache.get(real_sensor))
    return sensor_data

VIRTUAL_SENSORS = dict(DEFAULT_VIRTUAL_SENSORS)
VIRTUAL_SENSORS.update({'Antennas/{ant}/az' : _calc_azel, 'Antennas/{ant}/el' : _calc_azel})

#--------------------------------------------------------------------------------------------------
#--- Utility functions
#--------------------------------------------------------------------------------------------------

def get_single_value(group, name):
    """Return single value from attribute or dataset with given name in group.

    If `name` is an attribute of the HDF5 group `group`, it is returned,
    otherwise it is interpreted as an HDF5 dataset of `group` and the last value
    of `name` is returned. This is meant to retrieve static configuration values
    that potentially get set more than once during capture initialisation, but
    then does not change during actual capturing.

    Parameters
    ----------
    group : :class:`h5py.Group` object
        HDF5 group to query
    name : string
        Name of HDF5 attribute or dataset to query

    Returns
    -------
    value : object
        Attribute or last value of dataset

    """
    return group.attrs[name] if name in group.attrs else group[name].value[-1]

#--------------------------------------------------------------------------------------------------
#--- CLASS :  H5DataV2
#--------------------------------------------------------------------------------------------------

class H5DataV2(DataSet):
    """Load HDF5 format version 2 file produced by KAT-7 correlator.

    For more information on attributes, see the :class:`DataSet` docstring.

    Parameters
    ----------
    filename : string
        Name of HDF5 file
    ref_ant : string, optional
        Name of reference antenna, used to partition data set into scans
        (default is first antenna in use)
    time_offset : float, optional
        Offset to add to all correlator timestamps, in seconds

    Attributes
    ----------
    file : :class:`h5py.File` object
        Underlying HDF5 file, exposed via :mod:`h5py` interface

    """
    def __init__(self, filename, ref_ant='', time_offset=0.0):
        DataSet.__init__(self, filename, ref_ant, time_offset)

        # Load file
        self.file = f = h5py.File(filename, 'r')

        # Only continue if file is correct version and has been properly augmented
        self.version = f.attrs.get('version', '1.x')
        if not self.version.startswith('2.'):
            raise WrongVersion("Attempting to load version '%s' file with version 2 loader" % (self.version,))
        if not 'augment_ts' in f.attrs:
            raise BrokenFile('HDF5 file not augmented - please run k7_augment.py (provided by katcapture package)')

        # Load main HDF5 groups
        data_group, sensors_group, config_group = f['Data'], f['MetaData/Sensors'], f['MetaData/Configuration']
        markup_group = f['Markup']
        # Get observation script attributes, with defaults
        script_attrs = config_group['Observation'].attrs
        self.observer = script_attrs.get('script_observer', '')
        self.description = script_attrs.get('script_description', '')
        self.experiment_id = script_attrs.get('script_experiment_id', '')

        # ------ Extract timestamps ------

        self.dump_period = get_single_value(config_group['Correlator'], 'int_time')
        # Obtain visibility data and timestamps
        self._vis = data_group['correlator_data']
        self._timestamps = data_group['timestamps']
        num_dumps = len(self._timestamps)
        # Discard the last sample if the timestamp is a duplicate (caused by stop packet in k7_capture)
        num_dumps = (num_dumps - 1) if num_dumps > 1 and (self._timestamps[-1] == self._timestamps[-2]) else num_dumps
        # Estimate timestamps by assuming they are uniformly spaced (much quicker than loading them from file).
        # This is useful for the purpose of segmenting data set, where accurate timestamps are not that crucial.
        # The real timestamps are still loaded when the user explicitly asks for them.
        # Do quick test for uniform spacing of timestamps (necessary but not sufficient).
        if abs((self._timestamps[num_dumps - 1] - self._timestamps[0]) / self.dump_period + 1 - num_dumps) < 0.01:
            # Estimate the timestamps as being uniformly spaced
            data_timestamps = self._timestamps[0] + self.dump_period * np.arange(num_dumps)
        else:
            # Load the real timestamps instead and warn the user, as this is anomalous
            data_timestamps = self._timestamps[:num_dumps]
            expected_dumps = (self._timestamps[num_dumps - 1] - self._timestamps[0]) / self.dump_period + 1
            logger.warning(("Irregular timestamps detected in file '%s':"
                           "expected %.3f dumps based on dump period and start/end times, got %d instead") %
                           (filename, expected_dumps, num_dumps))
        # Move timestamps from start of each dump to the middle of the dump
        data_timestamps += 0.5 * self.dump_period + self.time_offset
        if data_timestamps[0] < 1e9:
            logger.warning("File '%s' has invalid first correlator timestamp (%f)" % (filename, data_timestamps[0],))
        self._time_keep = np.ones(num_dumps, dtype=np.bool)
        self.start_time = katpoint.Timestamp(data_timestamps[0] - 0.5 * self.dump_period)
        self.end_time = katpoint.Timestamp(data_timestamps[-1] + 0.5 * self.dump_period)

        # ------ Extract sensors ------

        # Populate sensor cache with all HDF5 datasets below sensor group that fit the description of a sensor
        cache = {}
        def register_sensor(name, obj):
            """A sensor is defined as a non-empty dataset with expected dtype."""
            if isinstance(obj, h5py.Dataset) and obj.shape != () and obj.dtype.names == ('timestamp','value','status'):
                # Rename pedestal sensors from the old regime to become sensors of the corresponding antenna
                name = ('Antennas/ant' + name[13:]) if name.startswith('Pedestals/ped') else name
                cache[name] = SensorData(obj, name)
        sensors_group.visititems(register_sensor)
        # Use estimated data timestamps for now, to speed up data segmentation
{'fraction_non_alphanumeric': 0.15081191663470145, 'fraction_numerical': 0.00869454034010996, 'mean_word_length': 2.4894044166852556, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '308767', 'n_tokens_mistral': 4472, 'n_tokens_neox': 4310, 'n_words': 1400}