A highly incomplete prototype for a TimeSeries class

example_timeseries.py

from collections import OrderedDict
import copy
import operator
from astropy.time import Time


class TimeSeries:
    def __init__(self, time, data=None, **kwargs):
        if not isinstance(time, Time):
            raise ValueError('time must be an astropy Time')
        self.time = time
        if data is None:
            self.data = OrderedDict()
        else:
            self.data = OrderedDict(data)

    def __len__(self):
        return len(self.time)

    def _operate(self, other, nddata_op, fallback_op):
        """
        Performs an operation on self and 'other'

        'other' can be a TimeSeries, a constant or compatible numpy array.
        If 'other' is a TimeSeries, all matching columns are operated on
        (non-matching columns are left alone). If 'other' is a constant or
        array, all columns are operated on by 'other'
        """
        if isinstance(other, TimeSeries):
            if len(self) != len(other):
                raise ValueError(
                    'input lengths [{:d} vs {:d}] do not match'.format(
                        len(self), len(other)
                    )
                )

            if not all(self.time == other.time):
                raise ValueError('input times do not match')

            result = copy.deepcopy(self)
            # arithmetic on matching columns
            data_set_self = set(self.data.keys())
            data_set_other = set(other.data.keys())
            for key in data_set_self.intersection(data_set_other):
                # assume NDDataRef subclass first
                if hasattr(self.data[key], nddata_op):
                    operator = getattr(self.data[key], nddata_op)
                    result.data[key] = operator(other.data[key])
                else:
                    result.data[key] = fallback_op(self.data[key], other.data[key])

        else:
            result = copy.deepcopy(self)
            for key in result.data.keys():
                result.data[key] = fallback_op(self.data[key], other)
        return result

    def __truediv__(self, other):
        """
        Divides the TimeSeries by 'other'.

        'other' can be a TimeSeries, a constant or compatible numpy array.
        If 'other' is a TimeSeries, all matching columns are divided
        (non-matching columns are left alone). If 'other' is a constant or
        array, all columns are divided by 'other'
        """
        return self._operate(other, 'divide', operator.truediv)

    def __mul__(self, other):
        return self._operate(other, 'multiply', operator.mul)

    def __add__(self, other):
        return self._operate(other, 'add', operator.add)

    def __sub__(self, other):
        return self._operate(other, 'sub', operator.sub)

    def __getitem__(self, aslice):
        copy_self = copy.deepcopy(self)
        copy_self.time = self.time[aslice]
        for key in copy_self.data:
            copy_self.data[key] = copy_self.data[key][aslice]
        return copy_self