time-split 0.1.0__py3-none-any.whl

time_split/__init__.py

@@ -0,0 +1,16 @@
+"""Create temporal k-folds for cross-validation with heterogeneous data."""
+
+import logging as _logging
+
+from ._frontend import log_split_progress, plot, split
+
+__all__ = [
+    "log_split_progress",
+    "plot",
+    "split",
+    "__version__",  # Make MyPy happy
+]
+
+__version__ = "0.1.0"
+
+_logging.getLogger(__name__).addHandler(_logging.NullHandler())

time_split/_backend/__init__.py

@@ -0,0 +1,4 @@
+from ._limits import expand_limits
+from ._splitter import DatetimeIndexSplitter
+
+__all__ = ["DatetimeIndexSplitter", "expand_limits"]

time_split/_backend/_available.py

@@ -0,0 +1,57 @@
+from typing import NamedTuple
+
+from pandas import DatetimeIndex, Timestamp, isna
+
+from .._support import handle_dask
+from ..types import DatetimeIterable, Flex
+from ._datetime_index_like import DatetimeIndexLike
+from ._limits import LimitsTuple, expand_limits
+
+
+class AvailableMetadata(NamedTuple):
+    """Output of :func:`process_available`."""
+
+    available_as_index: DatetimeIndexLike | None
+    limits: LimitsTuple
+    expanded_limits: LimitsTuple
+
+
+def process_available(available: DatetimeIterable, *, flex: Flex) -> AvailableMetadata:
+    """Process a user-given `available` argument.
+
+    Args:
+        available: Available data from user. May be ``None``
+        flex: Flex-argument. Determines how much (if at all) to expand limits.
+
+    Returns:
+        A tuple ``(available, limits)``. Note that `available` will be ``None``, it has not been iterated over. This
+        assures that iterables are not consumed unless needed.
+
+    Raises:
+        ValueError: For invalid `available` arguments.
+
+    """
+    if isinstance(available, DatetimeIndexLike):
+        # Avoids conversion, which may be expensive (or impossible on the current machine),
+        # and leverages additional non-Pandas types that do vectorized min/max.
+        limits = _compute_data_limits(available)
+        available_retval = available
+    else:
+        available_retval = DatetimeIndex(available)
+        limits = _compute_data_limits(available_retval)
+
+    min_dt, max_dt = limits
+    if min_dt == max_dt or isna(min_dt) or isna(max_dt):
+        raise ValueError("Not enough data in 'available'; at least two unique elements are required.")
+
+    if not (isinstance(min_dt, Timestamp) and isinstance(max_dt, Timestamp)):
+        # Enforce pandas types; we might have a numpy.datetime64
+        # (croniter 1.4.1 + numpy 1.25.2): croniter cannot handle numpy.datetime64
+        limits = Timestamp(min_dt), Timestamp(max_dt)
+
+    expanded_limits = expand_limits(limits, flex=flex)
+    return AvailableMetadata(available_retval, limits=limits, expanded_limits=expanded_limits)
+
+
+def _compute_data_limits(available: DatetimeIndexLike) -> LimitsTuple:
+    return handle_dask(available.min()), handle_dask(available.max())

time_split/_backend/_datetime_index_like.py

@@ -0,0 +1,19 @@
+from collections.abc import Iterable, Iterator
+from typing import Protocol, runtime_checkable
+
+from pandas import Series
+
+from ..types import DatetimeTypes
+
+
+@runtime_checkable
+class DatetimeIndexLike(Protocol, Iterable[DatetimeTypes]):
+    """A type that behaves (sort of) like a pandas DatetimeIndex."""
+
+    def __iter__(self) -> Iterator[DatetimeTypes]: ...
+
+    def value_counts(self) -> Series: ...
+
+    def min(self) -> DatetimeTypes: ...
+
+    def max(self) -> DatetimeTypes: ...

time_split/_backend/_limits.py

@@ -0,0 +1,151 @@
+from collections.abc import Iterable
+from itertools import starmap
+from typing import NamedTuple
+
+from pandas import Timedelta, Timestamp
+
+from ..settings import auto_flex
+from ..types import Flex, TimedeltaTypes
+
+LimitsTuple = tuple[Timestamp, Timestamp]
+LevelTuple = tuple[TimedeltaTypes, TimedeltaTypes, TimedeltaTypes]
+
+
+class _TimedeltaTuple(NamedTuple):
+    start_at: Timedelta
+    round_to: Timedelta
+    tolerance: Timedelta
+
+
+def expand_limits(limits: LimitsTuple, *, flex: Flex | LevelTuple | Iterable[LevelTuple] = "auto") -> LimitsTuple:
+    """Derive the `"real"` bounds of `limits`.
+
+    Args:
+        limits: A tuple ``(lo, hi)`` of timestamps.
+        flex: Flex arguments as described in the :ref:`User guide`. Also supports level-tuples
+            ``[(start_at, round_to, tolerance)...]``. Passing ``flex=[settings.auto_flex.day, settings.auto_flex.hour]``
+            is equivalent to ``flex='auto'``.
+
+    Returns:
+        Limits rounded according to the `flex`-argument.
+
+    Raises:
+        ValueError: For invalid limits.
+
+    Examples:
+        >>> from pandas import Timestamp
+        >>> limits = Timestamp("2019-05-11"), Timestamp("2019-05-11 22:05:30")
+
+        Basic usage.
+
+        >>> expand_limits(limits, flex="d")
+        (Timestamp('2019-05-11 00:00:00'), Timestamp('2019-05-12 00:00:00'))
+
+        You may specify a maximum "distance" that limits may be expanded.
+
+        >>> expand_limits(limits, flex="d<1h")
+        (Timestamp('2019-05-11 00:00:00'), Timestamp('2019-05-11 22:05:30'))
+
+        Limits will never be rounded in the "wrong" direction...
+
+        >>> limits = Timestamp("2019-05-11"), Timestamp("2019-05-11 11:05:30")
+        >>> expand_limits(limits, flex="d")
+        (Timestamp('2019-05-11 00:00:00'), Timestamp('2019-05-11 11:05:30'))
+
+        ...even if you make the tolerance large enough.
+
+        >>> expand_limits(limits, flex="d<14h")
+        (Timestamp('2019-05-11 00:00:00'), Timestamp('2019-05-11 11:05:30'))
+
+    """
+    if limits[0] >= limits[1]:
+        msg = f"Bad limits. Expected limits[1] > limits[0], but got {limits=}."
+        raise ValueError(msg)
+
+    if flex is False:
+        return limits
+
+    if flex is True or flex == "auto":
+        return _from_levels(limits)
+
+    if isinstance(flex, str):
+        round_to, _, tolerance = flex.partition("<")
+        level = _make_level(None, round_to=round_to.strip(), tolerance=tolerance.strip())
+        return _apply(limits, level=level)
+    if isinstance(flex, tuple):
+        return _from_levels(limits, levels=[_make_level(*flex)])
+
+    return _from_levels(limits, levels=starmap(_make_level, flex))
+
+
+def _from_levels(limits: LimitsTuple, *, levels: Iterable[_TimedeltaTuple] | None = None) -> LimitsTuple:
+    if levels is None:
+        levels = _levels_from_settings()
+    else:
+        levels = sorted(levels, reverse=True, key=lambda level: level[0])
+
+    lo, hi = limits
+    diff = hi - lo
+
+    for level in levels:
+        start_at, _round_to, _tolerance = level
+        if diff >= start_at:
+            return _apply(limits, level=level)
+
+    return limits
+
+
+def _apply(limits: LimitsTuple, *, level: _TimedeltaTuple) -> LimitsTuple:
+    lo, hi = limits
+
+    lo_floor = lo.floor(level.round_to)
+    if abs(lo_floor - lo) > level.tolerance:
+        return limits
+
+    hi_ceil = hi.ceil(level.round_to)
+    if abs(hi_ceil - hi) > level.tolerance:
+        return limits
+
+    if auto_flex.SANITY_CHECK:
+        if lo.round(level.round_to) != lo_floor:
+            return limits
+        if hi.round(level.round_to) != hi_ceil:
+            return limits
+
+    return lo_floor, hi_ceil
+
+
+def _levels_from_settings() -> list[_TimedeltaTuple]:
+    day, hour = _make_level(*auto_flex.day), _make_level(*auto_flex.hour)
+    if day.round_to != Timedelta(days=1):
+        msg = f"Invalid settings: {auto_flex.day=} must have round_to=1 day."
+        raise ValueError(msg)
+    if hour.round_to != Timedelta(hours=1):
+        msg = f"Invalid settings: {auto_flex.hour=} must have round_to=1 hour."
+        raise ValueError(msg)
+    return [day, hour]
+
+
+def _make_level(
+    start_at: TimedeltaTypes | None,
+    round_to: TimedeltaTypes,
+    tolerance: TimedeltaTypes,
+) -> _TimedeltaTuple:
+    start_at = Timedelta(start_at)  # Will be NaT if None, making all ifs False.
+    if start_at <= Timedelta(0):
+        raise ValueError(f"Bad {start_at=}; must be non-negative.")
+
+    try:
+        round_to = Timedelta(1, unit=round_to)
+    except ValueError as e:
+        raise ValueError(f"Got {round_to=}, which is not a valid frequency.") from e
+
+    tolerance = Timedelta(tolerance)
+    if tolerance <= Timedelta(0):
+        raise ValueError(f"Bad {tolerance=}; must be non-negative.")
+
+    if start_at < round_to:
+        raise ValueError(f"{start_at=} < {round_to=}")
+    if round_to < tolerance:
+        raise ValueError(f"{round_to=} < {tolerance=}")
+    return _TimedeltaTuple(start_at, round_to, tolerance)

time_split/_backend/_schedule.py

@@ -0,0 +1,66 @@
+from typing import NamedTuple, cast, get_args
+
+from pandas import DatetimeIndex, NaT, Timedelta, Timestamp, date_range
+
+from ..types import DatetimeIterable, Flex, Schedule, TimedeltaTypes
+from ._available import AvailableMetadata, process_available
+from ._limits import LimitsTuple
+
+NO_LIMITS = cast(LimitsTuple, (NaT, NaT))
+
+
+class MaterializedSchedule(NamedTuple):
+    """An explicit schedule."""
+
+    schedule: DatetimeIndex
+    available_metadata: AvailableMetadata
+
+
+def materialize_schedule(
+    schedule: Schedule, flex: Flex, *, available: DatetimeIterable | None = None
+) -> MaterializedSchedule:
+    """Materialize user schedule based on available data."""
+    if available is None:
+        try:
+            return MaterializedSchedule(
+                DatetimeIndex(schedule),
+                available_metadata=AvailableMetadata(None, NO_LIMITS, NO_LIMITS),
+            )
+        except TypeError as e:
+            raise ValueError("Schedule must be explicit when not bounded by an available range.") from e
+
+    available_metadata = process_available(available, flex=flex)
+    min_dt, max_dt = available_metadata.expanded_limits
+
+    if isinstance(schedule, str) and _cron_like(schedule):
+        schedule = _handle_cron(schedule, min_dt, max_dt)
+    elif isinstance(schedule, get_args(TimedeltaTypes)):
+        schedule = _from_timedelta(schedule, available_metadata.expanded_limits)
+    else:
+        if not isinstance(schedule, DatetimeIndex):
+            schedule = DatetimeIndex(schedule)
+
+        schedule = schedule.tz_localize(min_dt.tz) if schedule.tz is None else schedule.tz_convert(min_dt.tz)
+        schedule = schedule[(min_dt <= schedule) & (schedule <= max_dt)]
+    return MaterializedSchedule(schedule, available_metadata=available_metadata)
+
+
+def _cron_like(schedule: str) -> bool:
+    # CroniterBadCronError: Exactly 5 or 6 columns has to be specified for iterator expression.
+    return len(schedule.split()) > 4 or schedule[0] == "@"  # noqa: PLR2004
+
+
+def _from_timedelta(schedule: TimedeltaTypes, limits: LimitsTuple) -> DatetimeIndex:
+    timedelta = Timedelta(schedule)
+    if timedelta <= Timedelta(0):
+        raise ValueError(f"unbounded {schedule=} must be greater than zero.")
+    return date_range(*limits, freq=timedelta, inclusive="both")
+
+
+def _handle_cron(expr: str, min_dt: Timestamp, max_dt: Timestamp) -> DatetimeIndex:
+    try:
+        from croniter import croniter_range
+
+        return DatetimeIndex(croniter_range(min_dt, max_dt, expr))
+    except ModuleNotFoundError as e:
+        raise ModuleNotFoundError(f"Install 'croniter' to parse cron expressions such such as '{expr}'.") from e

time_split/_backend/_span.py

@@ -0,0 +1,128 @@
+from collections.abc import Callable
+from datetime import timedelta
+from typing import Literal, TypeAlias
+
+from pandas import DatetimeIndex, Timedelta, Timestamp
+
+from ..types import Span
+from ._limits import LimitsTuple
+from ._schedule import NO_LIMITS
+
+StrictSpan: TypeAlias = int | Timedelta | Literal["all"]
+SpanArgumentName = Literal["before", "after"]
+FuncType = Callable[[int], Timestamp | None]
+
+
+class InvalidSpanError(ValueError):
+    """Raised when offsets values for 'before' and 'after' arguments are invalid."""
+
+    def __init__(
+        self,
+        span: Span,
+        *,
+        name: SpanArgumentName,
+        reason: str = "must be greater than zero",
+    ) -> None:
+        super().__init__(f"Bad '{name}'-argument; offset {span!r} {reason}.")
+
+
+def _to_timedelta(span: str | timedelta | Timedelta, *, name: SpanArgumentName) -> Timedelta:
+    retval = Timedelta(span)
+    if retval <= Timedelta(0):
+        raise InvalidSpanError(span, name=name)
+    return retval
+
+
+def to_strict_span(span: Span, *, name: SpanArgumentName) -> StrictSpan:
+    """Convert `span` to a strict span.
+
+    Args:
+        span: User span.
+        name: Name of user input. User for error reporting.
+
+    Returns:
+        A valid span with reduced types.
+
+    Raises:
+        TypeError: For bad `span` types.
+        InvalidSpanError: For invalid `span` arguments.
+
+    """
+    if span == "all":
+        return "all"
+
+    if isinstance(span, (str, timedelta, Timedelta)):
+        return _to_timedelta(span, name=name)
+
+    if span is False or isinstance(span, (int, float)):
+        retval = int(span)
+        if retval <= 0:
+            raise InvalidSpanError(retval, name=name)
+        return retval
+
+    raise TypeError(f"cannot parse '{name}'-argument of type {type(span).__name__}: {span!r}.")
+
+
+class OffsetCalculator:
+    """Utility class for computing before/after offsets from the `schedule`."""
+
+    def __init__(
+        self,
+        span: Span,
+        schedule: DatetimeIndex,
+        limits: LimitsTuple,
+        *,
+        name: SpanArgumentName,
+    ) -> None:
+        span = to_strict_span(span, name=name)
+        is_before = name == "before"
+
+        self.func: FuncType
+        if span == "all":
+            if limits == NO_LIMITS:
+                raise InvalidSpanError(
+                    span,
+                    name=name,
+                    reason="requires available data to bound the schedule",
+                )
+            self.func = self._make_all(is_before, limits)
+        elif isinstance(span, int):
+            self.func = self._make_int(is_before, span, schedule)
+        else:  # Timedelta
+            self.func = self._make_timedelta(is_before, span, schedule)
+
+    @staticmethod
+    def _make_all(is_before: bool, limits: LimitsTuple) -> FuncType:
+        retval = Timestamp(limits[0 if is_before else 1])
+
+        def func(_: int) -> Timestamp | None:
+            return retval
+
+        return func
+
+    @staticmethod
+    def _make_int(is_before: bool, offset: int, schedule: DatetimeIndex) -> FuncType:
+        def func(i: int) -> Timestamp | None:
+            i = i - offset if is_before else i + offset
+            if 0 <= i < len(schedule):
+                return schedule[i]
+            else:
+                return None
+
+        return func
+
+    @staticmethod
+    def _make_timedelta(is_before: bool, offset: Timedelta, schedule: DatetimeIndex) -> FuncType:
+        if is_before:
+            schedule = schedule - offset
+        else:
+            schedule = schedule + offset
+
+        def func(i: int) -> Timestamp | None:
+            return schedule[i]
+
+        return func
+
+    def __call__(self, i: int) -> Timestamp | None:
+        """Get start/end timestamp."""
+        return self.func(i)

time_split/_backend/_splitter.py

@@ -0,0 +1,162 @@
+import logging
+from dataclasses import asdict, dataclass
+from typing import cast, get_args
+
+from pandas import Timedelta, Timestamp
+from rics.misc import format_kwargs
+from rics.performance import format_seconds
+
+from ..settings import misc
+from ..types import (
+    DatetimeIndexSplitterKwargs,
+    DatetimeIterable,
+    DatetimeSplitBounds,
+    DatetimeSplits,
+    Flex,
+    Schedule,
+    Span,
+    TimedeltaTypes,
+)
+from ._limits import LimitsTuple
+from ._schedule import MaterializedSchedule, materialize_schedule
+from ._span import OffsetCalculator, to_strict_span
+
+LOGGER = logging.getLogger("time_split")
+
+
+@dataclass(frozen=True)
+class DatetimeIndexSplitter:
+    """Backend interface for splitting user data."""
+
+    schedule: Schedule
+    before: Span
+    after: Span
+    step: int
+    n_splits: int | None
+    flex: Flex
+
+    def get_splits(self, available: DatetimeIterable | None = None) -> DatetimeSplits:
+        """Compute a split of given user data."""
+        ms = self._materialize_schedule(available)
+        self._log_expansion(ms.available_metadata.limits, expanded=ms.available_metadata.expanded_limits)
+        return self._make_bounds_list(ms)
+
+    def get_plot_data(self, available: DatetimeIterable | None = None) -> tuple[DatetimeSplits, MaterializedSchedule]:
+        """Returns additional data needed to visualize folds."""
+        ms = self._materialize_schedule(available)
+        splits = self._make_bounds_list(ms)
+        return splits, ms
+
+    def _materialize_schedule(self, available: DatetimeIterable | None = None) -> MaterializedSchedule:
+        ms = materialize_schedule(self.schedule, self.flex, available=available)
+        if not ms.schedule.sort_values().equals(ms.schedule):
+            raise ValueError(f"schedule must be sorted in ascending order; schedule={self.schedule!r} is not valid.")
+
+        types = get_args(TimedeltaTypes)
+        if (
+            misc.snap_to_end
+            and self.after != "all"
+            and isinstance(self.schedule, types)
+            and isinstance(self.after, types)
+        ):
+            ms = self._snap_to_end(ms)
+
+        return ms
+
+    def _snap_to_end(self, ms: MaterializedSchedule) -> MaterializedSchedule:
+        schedule_frequency = ms.schedule.freq
+        if schedule_frequency is None:
+            return ms
+
+        from_end = ms.available_metadata.expanded_limits[1] - (ms.schedule[-1] + Timedelta(self.after))
+        from_end = from_end.floor(schedule_frequency.base)
+
+        return ms._replace(schedule=ms.schedule + from_end)
+
+    def _make_bounds_list(self, ms: MaterializedSchedule) -> DatetimeSplits:
+        get_start = OffsetCalculator(
+            self.before,
+            ms.schedule,
+            ms.available_metadata.expanded_limits,
+            name="before",
+        )
+        get_end = OffsetCalculator(self.after, ms.schedule, ms.available_metadata.expanded_limits, name="after")
+
+        min_start, max_end = ms.available_metadata.expanded_limits
+
+        retval = []
+        for i, mid in enumerate(ms.schedule):
+            start, end = get_start(i), get_end(i)
+            if start is None or start < min_start or start >= mid:
+                continue
+            if end is None or end > max_end or end <= mid:
+                continue
+            retval.append(DatetimeSplitBounds(start, mid, end))
+
+        if not retval:
+            limits_info = f"limits={tuple(map(str, ms.available_metadata.limits))} and "
+            msg = f"No valid splits with {limits_info}split params: ({format_kwargs(self.as_dict())})"
+            raise ValueError(msg)
+
+        return self._filter(retval)
+
+    def _filter(self, splits: DatetimeSplits) -> DatetimeSplits:
+        """Apply splitting arguments.
+
+        Args:
+            splits: Splits to filter.
+
+        Returns:
+            Filtered splits.
+
+        """
+        if self.step != 1:
+            step = abs(self.step)
+            splits = [s for i, s in enumerate(reversed(splits)) if i % step == 0]
+            splits.reverse()
+
+        if self.n_splits:
+            splits = splits[-self.n_splits :]
+
+        if self.step < 0:
+            splits.reverse()
+
+        return splits
+
+    def _log_expansion(self, original: LimitsTuple, *, expanded: LimitsTuple) -> None:
+        if original == expanded:
+            return
+
+        if not LOGGER.isEnabledFor(logging.INFO):
+            return
+
+        def stringify(old: Timestamp, *, new: Timestamp) -> str:
+            from .._frontend._to_string import _PrettyTimestamp
+
+            retval = f"{old} -> "
+            if old == new:
+                return retval + "<no change>"
+            diff = (new - old).total_seconds()
+            return retval + f"{_PrettyTimestamp(new).auto} ({'+' if diff > 0 else '-'}{format_seconds(abs(diff))})"
+
+        LOGGER.info(
+            f"Available data limits have been expanded (since flex={self.flex!r}):\n"
+            f"  start: {stringify(original[0], new=expanded[0])}\n"
+            f"    end: {stringify(original[1], new=expanded[1])}"
+        )
+
+    def __post_init__(self) -> None:
+        # Verify n_splits
+        if self.n_splits is not None and self.n_splits < 1:
+            raise ValueError(f"Expected n_splits >= 1, but got n_splits={self.n_splits!r}.")
+
+        # Verify before/after
+        to_strict_span(self.before, name="before")
+        to_strict_span(self.after, name="after")
+
+        if self.step == 0:
+            raise ValueError(f"Bad argument step={self.step}; must be a non-zero integer.")
+
+    def as_dict(self) -> DatetimeIndexSplitterKwargs:
+        """Returns the splitter as a ``dict``."""
+        return cast(DatetimeIndexSplitterKwargs, asdict(self))

time_split/_docstrings.py

@@ -0,0 +1,36 @@
+from typing import TypeVar
+
+FuncType = TypeVar("FuncType")  # No proper way to type this?
+
+
+def add_docstrings(**kwargs: str) -> None:
+    """Add a documented argument to the ``@docs`` decorator."""
+    _DOCSTRINGS.update(kwargs)
+
+
+def docs(function: FuncType) -> FuncType:
+    """Add info the docstring of a function."""
+    assert function.__doc__ is not None  # noqa: S101
+    function.__doc__ = function.__doc__.format(**_DOCSTRINGS)
+    return function
+
+
+_OFFSET = "pandas :ref:`offset alias <pandas:timeseries.offset_aliases>`"
+_SPAN = (
+    f"Range {{arg}} `schedule` timestamps. Either a {_OFFSET}, an integer (`schedule`-based offsets), "
+    f"or `'all'` (requires `available` data)."
+)
+_DOCSTRINGS = {
+    "schedule": f"A collection of timestamps, a {_OFFSET}, or a cron expression.",
+    "before": _SPAN.format(arg="before"),
+    "after": _SPAN.format(arg="after"),
+    "step": "Select a subset of folds, preferring folds later in the schedule.",
+    "n_splits": "Maximum number of folds, preferring folds later in the schedule.",
+    "available": "Binds `schedule` to a range.",
+    "flex": f'A {_OFFSET} used to expand `available` data to its likely `"true"` limits. Pass ``False`` to disable.',
+    "USER_GUIDE": (
+        "For more information about the `schedule`, `before/after` and `flex`-arguments" ", see the :ref:`User guide`."
+    ),
+    "OFFSET": _OFFSET,
+    "DatetimeIndexSplitterKwargs": "See :func:`~time_split.split`. The `available` keyword is managed by the integration.",
+}

time_split/_frontend/__init__.py

@@ -0,0 +1,13 @@
+from ._plot import plot
+from ._progress import log_split_progress
+from ._split import split
+from ._to_string import to_string
+from ._weight import fold_weight
+
+__all__ = [
+    "fold_weight",
+    "log_split_progress",
+    "plot",
+    "split",
+    "to_string",
+]