Source code for irfpy.util.span

''' Module related to span (interval).

.. codeauthor:: Yoshifumi Futaana

The code depends on ``sympy`` package.  Version 1.5.1 is recommended (2020-03-27).
'''
import sympy

from irfpy.util import timeinterval
from irfpy.util import version as _v

_v.recommends("1.5.1", sympy.__version__, "sympy")

def _get_empty_set():
    try:
        e = sympy.EmptySet()
    except TypeError:
        e = sympy.EmptySet

    return e


def array2interval(true_false_array, sorted_array=None):
    ''' Two arrays are converted to ``interval.IntervalSet``.

    :param true_false_array: An array that contains ``true`` and ``false``.
    :param sorted_array: An array that contains number. Each component
        should be comparable each other. The array should be sorted from small to large.
        In case of ``None``, index is used.

    >>> tfarr = [False, False, False, True, True, True, True, False, True, True, True]
    >>> numarr = [1.5, 2.8, 2.9, 3.5, 3.8, 3.9, 5.1, 5.8, 5.9, 6.1, 6.3]

    In the above case, 3.5 (index 3) to 5.1 (index 6) and 5.9 to 6.3 (last 3 elements)
    is true, so that

    >>> intrv = array2interval(tfarr, numarr)
    >>> print((3.4 in intrv), (3.5 in intrv))
    False True
    >>> print((5.1 in intrv), (5.2 in intrv))
    True False
    >>> print((5.8 in intrv), (5.9 in intrv))
    False True
    >>> print((6.3 in intrv), (6.4 in intrv))
    True False

    If only ``true_false_array`` is given, the index (3-6 and 8-10) is used.

    >>> intrv = array2interval(tfarr)
    >>> print((2 in intrv), (3 in intrv))
    False True
    >>> print((6 in intrv), (7 in intrv))
    True False
    >>> print((7 in intrv), (8 in intrv))
    False True
    >>> print((10 in intrv), (11 in intrv))
    True False

    '''
    if sorted_array is None:
        sorted_array = list(range(len(true_false_array)))

    if len(true_false_array) != len(sorted_array):
        raise ValueError('Given arrays have different length. (%d and %d)' % (len(true_false_array), len(sorted_array)))

    n = len(true_false_array)

    idx0 = None

    prevtf = true_false_array[0]
    if prevtf:
        idx0 = 0

    intervalset = _get_empty_set()

    for i in range(1, n):
        nowtf = true_false_array[i]
        if nowtf:
            if prevtf:
                pass
            else:
                idx0 = i
        else:
            if prevtf:
                intvl = sympy.Interval(sorted_array[idx0], sorted_array[i - 1])
                intervalset += intvl
                idx0 = None
            else:
                pass
        prevtf = nowtf

    if idx0 is not None:  # Still in the true span
        intvl = sympy.Interval(sorted_array[idx0], sorted_array[-1])
        intervalset += intvl

    return intervalset


def timearray2interval(true_false_array, time_array):
    """

    :param true_false_array: List of "True" or "False"
    :type true_false_array: list of boolean
    :param time_array: List of ``datetime.datetime``.
    :type time_array: list of ``datetime.datetime``
    :return: The time interval corresponding to the true_false_array.
    :rtype: :class:`irfpy.util.timeinterval.timeinterval`
    """

    if len(true_false_array) != len(time_array):
        raise ValueError('Given arrays have different length. (%d and %d)' % (len(true_false_array), len(time_array)))

    n = len(true_false_array)

    idx0 = None

    prevtf = true_false_array[0]
    if prevtf:
        idx0 = 0

    intervalset = timeinterval.empty_interval()

    for i in range(1, n):
        nowtf = true_false_array[i]
        if nowtf:
            if prevtf:
                pass
            else:
                idx0 = i
        else:
            if prevtf:
                intvl = timeinterval.timeinterval(time_array[idx0], time_array[i - 1])
                intervalset += intvl
                idx0 = None
            else:
                pass
        prevtf = nowtf

    if idx0 is not None:  # Still in the true span
        intvl = timeinterval.timeinterval(time_array[idx0], time_array[-1])
        intervalset += intvl

    return intervalset